Control of the sketched least-squares cross term

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.SampleLem40pf

/-- Index identification `Fin (2^k) ⊕ Fin (2^k) ≃ Fin (2^(k+1))`: the first
summand is the first half of the indices, the second summand the second half. -/
def hadamardIndexEquiv (k : ℕ) : Fin (2 ^ k) ⊕ Fin (2 ^ k) ≃ Fin (2 ^ (k + 1)) :=
  finSumFinEquiv.trans (finCongr (by ring))

/-- The unnormalized Hadamard matrix `H̃_{2^k}` from the source recursion
`H̃_{2m} = [[H̃_m, H̃_m], [H̃_m, -H̃_m]]`. The `k = 0` value `[1]` is only a
construction device: at `k = 1` it yields the source base case
`H̃_2 = [[+1, +1], [+1, -1]]`; only sizes `2^(k+1) ≥ 2` are used below. -/
noncomputable def unnormalizedHadamard : (k : ℕ) → Matrix (Fin (2 ^ k)) (Fin (2 ^ k)) ℝ
  | 0 => fun _ _ => 1
  | k + 1 => Matrix.reindex (hadamardIndexEquiv k) (hadamardIndexEquiv k)
      (Matrix.fromBlocks (unnormalizedHadamard k) (unnormalizedHadamard k)
        (unnormalizedHadamard k) (-unnormalizedHadamard k))

/-- The normalized Hadamard transform `H_n = (1/√n) H̃_n`, `n = 2^(k+1)`. -/
noncomputable def normalizedHadamard (k : ℕ) :
    Matrix (Fin (2 ^ (k + 1))) (Fin (2 ^ (k + 1))) ℝ :=
  (1 / Real.sqrt ((2 ^ (k + 1) : ℕ) : ℝ)) • unnormalizedHadamard (k + 1)

/-- Law of a diagonal entry of `D`: `+1` w.p. `1/2`, `-1` w.p. `1/2`. -/
noncomputable def signLaw : Measure ℝ :=
  (1 / 2 : ℝ≥0∞) • Measure.dirac (1 : ℝ) + (1 / 2 : ℝ≥0∞) • Measure.dirac (-1 : ℝ)

/-- Uniform law on the index set `{1,…,n}` (here `Fin n`): each index has mass `1/n`.
This is the source's "select uniformly at random an integer from {1,2,…,n}". -/
noncomputable def uniformIndexLaw (n : ℕ) : Measure (Fin n) :=
  ((n : ℝ≥0∞)⁻¹) • Measure.count

/-- The `n × r` sampling-and-rescaling matrix of RandLeastSquares: column `t` is
`√(n/r) e_i` where `i = idx t` is the index selected in trial `t`. -/
noncomputable def samplingMatrix {n r : ℕ} (idx : Fin r → Fin n) : Matrix (Fin n) (Fin r) ℝ :=
  fun i t => if idx t = i then Real.sqrt ((n : ℝ) / (r : ℝ)) else 0

/-- Randomness of RandLeastSquares: diagonal signs of `D` are `±1` w.p. `1/2`
each; the `r` sampled indices are uniform on `{1,…,n}` (i.i.d. trials with
replacement); all draws are mutually independent (indices via real casts). -/
structure RandLeastSquaresRandomness {Ω : Type*} [MeasurableSpace Ω] {n r : ℕ}
    (P : Measure Ω) (sgn : Ω → Fin n → ℝ) (idx : Ω → Fin r → Fin n) : Prop where
  isProbabilityMeasure : IsProbabilityMeasure P
  sign_law : ∀ i, ProbabilityTheory.HasLaw (fun ω => sgn ω i) signLaw P
  index_law : ∀ t, ProbabilityTheory.HasLaw (fun ω => idx ω t) (uniformIndexLaw n) P
  indep : ProbabilityTheory.iIndepFun
    (fun (j : Fin n ⊕ Fin r) (ω : Ω) =>
      Sum.elim (fun i => sgn ω i) (fun t => (((idx ω t : Fin n) : ℕ) : ℝ)) j) P

/-- Eqn. (eqn:lem:HU_eqn2) for `M = H D U`:
`‖M_{i*}‖_2^2 ≤ 2 d ln(40 n d) / n` for all `i`. -/
def HURowNormBound {n d : ℕ} (M : Matrix (Fin n) (Fin d) ℝ) : Prop :=
  ∀ i, ∑ j, (M i j) ^ 2 ≤ 2 * (d : ℝ) * Real.log (40 * (n : ℝ) * (d : ℝ)) / (n : ℝ)

/-- `b^⊥ = U_A^⊥ U_A^⊥ᵀ b = b - U_A U_Aᵀ b`. -/
def residual {n d : ℕ} (UA : Matrix (Fin n) (Fin d) ℝ) (b : Fin n → ℝ) : Fin n → ℝ :=
  b - UA *ᵥ (UAᵀ *ᵥ b)

/-- `(Sᵀ H D U_A)ᵀ Sᵀ H D b^⊥`. -/
noncomputable def sketchedCrossProduct {n d r : ℕ} (H : Matrix (Fin n) (Fin n) ℝ)
    (sgn : Fin n → ℝ) (idx : Fin r → Fin n) (UA : Matrix (Fin n) (Fin d) ℝ)
    (bperp : Fin n → ℝ) : Fin d → ℝ :=
  ((samplingMatrix idx)ᵀ * H * Matrix.diagonal sgn * UA)ᵀ *ᵥ
    (((samplingMatrix idx)ᵀ * H * Matrix.diagonal sgn) *ᵥ bperp)

/-- Lemma lem:sample_lem40pf (Drineas–Mahoney 2017). With `n = 2^(k+1)`,
`H` the normalized Hadamard transform, `D = diag(sgn)`, `S` the sampling matrix,
`U_A` an orthonormal basis (left singular vectors) of the column space of `A`,
`b^⊥ = U_A^⊥U_A^⊥ᵀ b`, `Z = ‖b^⊥‖_2`: if `r ≥ 40 d ln(40nd)/ε`, then conditioned on
Eqn. (eqn:lem:HU_eqn2) for `HDU_A`, the event
`‖(SᵀHDU_A)ᵀ SᵀHD b^⊥‖_2^2 ≤ ε Z^2/2` has probability at least `.9`, written as
`0.9 · P(HU) ≤ P(HU ∩ E)`. -/
theorem sample_lem40pf {Ω : Type*} [MeasurableSpace Ω] (P : Measure Ω)
    (k d r : ℕ) (ε : ℝ) (hε0 : 0 < ε) (hε1 : ε < 1)
    (A : Matrix (Fin (2 ^ (k + 1))) (Fin d) ℝ) (b : Fin (2 ^ (k + 1)) → ℝ)
    (UA : Matrix (Fin (2 ^ (k + 1))) (Fin d) ℝ)
    (hUA : OrthonormalColumns UA)
    (hUAA : matrixColumnSpace UA = matrixColumnSpace A)
    (sgn : Ω → Fin (2 ^ (k + 1)) → ℝ) (idx : Ω → Fin r → Fin (2 ^ (k + 1)))
    (hrand : RandLeastSquaresRandomness P sgn idx)
    (hr : 40 * (d : ℝ) * Real.log (40 * ((2 ^ (k + 1) : ℕ) : ℝ) * (d : ℝ)) / ε ≤ (r : ℝ)) :
    (9 / 10 : ℝ≥0∞) *
        P {ω | HURowNormBound (normalizedHadamard k * Matrix.diagonal (sgn ω) * UA)} ≤
      P {ω | HURowNormBound (normalizedHadamard k * Matrix.diagonal (sgn ω) * UA) ∧
        l2Norm (sketchedCrossProduct (normalizedHadamard k) (sgn ω) (idx ω) UA
            (residual UA b)) ^ 2 ≤
          ε * (l2Norm (residual UA b)) ^ 2 / 2} := by
  sorry

end RNLAPhaseB.DM2017.SampleLem40pf