Entrywise expectation and variance of sampled multiplication

import RNLAPhaseB.B3Base
import RNLAPhaseB.Complexity

open MeasureTheory
open scoped Matrix
open scoped Matrix.Norms.L2Operator
open scoped BigOperators
open scoped ENNReal

namespace RNLAPhaseB.DM2017.MultExpVar

/-- Sampling step of RandMatrixMultiply: the `c` column/row indices `i_1, …, i_c`
are drawn independently and with replacement from `{1,…,n}`, with
`P[i_t = k] = p_k` for every trial `t` and every index `k`. -/
def IIDIndexSampling {Ω : Type*} [MeasurableSpace Ω] (P : Measure Ω) {n c : ℕ}
    (prob : Fin n → ℝ) (idx : Fin c → Ω → Fin n) : Prop :=
  (∀ t, Measurable (idx t)) ∧
  (∀ t k, P {ω | idx t ω = k} = ENNReal.ofReal (prob k)) ∧
  ProbabilityTheory.iIndepFun idx P

/-- The matrix `C` of RandMatrixMultiply: its `t`-th column is
`A_{*i_t} / sqrt(c p_{i_t})`. -/
noncomputable def randMatMulC {m n c : ℕ} (A : Matrix (Fin m) (Fin n) ℝ)
    (prob : Fin n → ℝ) (idx : Fin c → Fin n) : Matrix (Fin m) (Fin c) ℝ :=
  fun i t => A i (idx t) / Real.sqrt ((c : ℝ) * prob (idx t))

/-- The matrix `R` of RandMatrixMultiply: its `t`-th row is
`B_{i_t *} / sqrt(c p_{i_t})`. -/
noncomputable def randMatMulR {n p c : ℕ} (B : Matrix (Fin n) (Fin p) ℝ)
    (prob : Fin n → ℝ) (idx : Fin c → Fin n) : Matrix (Fin c) (Fin p) ℝ :=
  fun t j => B (idx t) j / Real.sqrt ((c : ℝ) * prob (idx t))

/-- dm2017 Lemma (lem:multexpvar). Let `C` and `R` be constructed as in the
RandMatrixMultiply algorithm (input `A ∈ ℝ^{m×n}`, `B ∈ ℝ^{n×p}`, an integer `c`
with `1 ≤ c ≤ n`, and probabilities `p_k ≥ 0` with `∑ p_k = 1`; `c` indices drawn
i.i.d. with `P[i_t = k] = p_k`, `C_{*t} = A_{*i_t}/sqrt(c p_{i_t})`,
`R_{t*} = B_{i_t*}/sqrt(c p_{i_t})`). Then for every entry `(i,j)`,
`E[(CR)_{ij}] = (AB)_{ij}` and
`Var[(CR)_{ij}] ≤ (1/c) ∑_k A_{ik}^2 B_{kj}^2 / p_k`. -/
theorem randMatrixMultiply_entry_expectation_variance
    {Ω : Type*} [MeasurableSpace Ω] (P : Measure Ω) [IsProbabilityMeasure P]
    {m n p c : ℕ} (A : Matrix (Fin m) (Fin n) ℝ) (B : Matrix (Fin n) (Fin p) ℝ)
    (prob : Fin n → ℝ) (hprob_nonneg : ∀ k, 0 ≤ prob k) (hprob_sum : ∑ k, prob k = 1)
    (hc_pos : 1 ≤ c) (hc_le : c ≤ n)
    (idx : Fin c → Ω → Fin n) (hidx : IIDIndexSampling P prob idx) :
    ∀ (i : Fin m) (j : Fin p),
      (∫ ω, (randMatMulC A prob (fun t => idx t ω) *
          randMatMulR B prob (fun t => idx t ω)) i j ∂P) = (A * B) i j ∧
      ProbabilityTheory.variance
          (fun ω => (randMatMulC A prob (fun t => idx t ω) *
            randMatMulR B prob (fun t => idx t ω)) i j) P
        ≤ (1 / (c : ℝ)) * ∑ k : Fin n, A i k ^ 2 * B k j ^ 2 / prob k := by
  sorry

end RNLAPhaseB.DM2017.MultExpVar

/-!
SOURCE WARNING sw_late_ea7dd3bfeb5b4cc3: dm2017::lem:multexpvar expectation clause E[(CR)_ij]=(AB)_ij under the RandMatrixMultiply input condition p_k ≥ 0. The algorithm explicitly allows zero sampling probabilities: its input requires only p_k ≥ 0 with sum 1. It draws i_t with P[i_t = k] = p_k. The lemma applies to C and R 'constructed as described in the RandMatrixMultiply algorithm'. The proof's expectation step cancels p_k against 1/(c p_k), which is valid only when p_k > 0. The proof then asserts unbiasedness 'regardless of the choice of the sampling probabilities'. If p_k = 0 for some k, that index is drawn with probability zero, so its term A_ik B_kj never contributes to E[(CR)_ij]. The expectation then equals the sum of A_ik B_kj over k with p_k > 0. The claimed equality with (AB)_ij therefore fails whenever some p_k = 0 has A_ik B_kj ≠ 0. So the source's stated hypotheses (p_k ≥ 0) are inconsistent with its stated conclusion; the sources are internally contradictory. The candidate reproduces these hypotheses exactly: hprob_nonneg : ∀ k, 0 ≤ prob k and hprob_sum. It does not add a positivity assumption. The sampling law P{idx t = k} = ofReal(prob k) and the definitions of C and R match the algorithm. The expectation clause is the source's claimed conclusion verbatim. The null event {idx t = k} with p_k = 0 contributes nothing to the integral, independent of Lean's x/0 convention. The candidate therefore inherits exactly the source defect, not a defect of its own. No mathematical change is needed to be faithful. Adding p_k > 0 would be a repair of the source, which is not authorized. The variance clause's RHS term A_ik²B_kj²/p_k is undefined in the source when p_k = 0. That concerns a different clause and is not part of this finding, so it creates no additional objection to the proposal.
The source is stated as written; no mathematical declaration was changed.
-/