Tail interaction with the sampled leading subspace

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

/-- The normalized Sylvester-Hadamard matrix `H_n` with `n = 2^p`. Its `(i,j)` entry is
`n^(-1/2) * (-1)^(number of common one-bits of i and j)`. -/
noncomputable def normalizedHadamard (p : ℕ) : Matrix (Fin (2 ^ p)) (Fin (2 ^ p)) ℝ :=
  fun i j => (Real.sqrt ((2 : ℝ) ^ p))⁻¹ *
    (-1 : ℝ) ^ ((Finset.univ.filter
      (fun b : Fin p => (i.val.testBit b.val && j.val.testBit b.val) = true)).card)

/-- The first `k` columns of a matrix (`V_k` from `V`). -/
def leadingColumns {r n : ℕ} (k : ℕ) (hk : k ≤ n) (V : Matrix (Fin r) (Fin n) ℝ) :
    Matrix (Fin r) (Fin k) ℝ :=
  V.submatrix id (Fin.castLE hk)

/-- A full SVD `A = U Sig Vᵀ`: `U`, `V` orthogonal, `Sig` rectangular diagonal carrying the
nonnegative, nonincreasing singular values `s`. -/
structure IsSVD {m n : ℕ} (A : Matrix (Fin m) (Fin n) ℝ) (U : Matrix (Fin m) (Fin m) ℝ)
    (Sig : Matrix (Fin m) (Fin n) ℝ) (V : Matrix (Fin n) (Fin n) ℝ) (s : ℕ → ℝ) : Prop where
  U_orth : Uᵀ * U = 1
  V_orth : Vᵀ * V = 1
  sig_entries : ∀ i j, Sig i j = if i.val = j.val then s i.val else 0
  s_nonneg : ∀ i, 0 ≤ s i
  s_antitone : Antitone s
  factor : A = U * Sig * Vᵀ

/-- `A_{k,⊥} = A - A_k`, where `A_k = A V_k V_kᵀ` (equal to `U_k Σ_k V_kᵀ` under the SVD). -/
noncomputable def kPerp {m n k : ℕ} (A : Matrix (Fin m) (Fin n) ℝ)
    (Vk : Matrix (Fin n) (Fin k) ℝ) : Matrix (Fin m) (Fin n) ℝ :=
  A - A * Vk * Vkᵀ

/-- Squared Frobenius norm. -/
def frobSq {a b : ℕ} (M : Matrix (Fin a) (Fin b) ℝ) : ℝ :=
  ∑ i, ∑ j, (M i j) ^ 2

/-- The `n × c` uniform sampling-and-rescaling matrix: column `t` is `sqrt(n/c) e_{idx t}`. -/
noncomputable def uniformSamplingMatrix (n c : ℕ) (idx : Fin c → Fin n) :
    Matrix (Fin n) (Fin c) ℝ :=
  fun i t => if idx t = i then Real.sqrt ((n : ℝ) / c) else 0

/-- The `c` sampled indices are i.i.d. uniform on `Fin n` (sampling with replacement). -/
structure UniformSamplingIndices {Ω : Type*} [MeasurableSpace Ω] (n c : ℕ)
    (P : Measure Ω) (idx : Ω → Fin c → Fin n) : Prop where
  isProbabilityMeasure : IsProbabilityMeasure P
  hasLaw : ∀ t, ProbabilityTheory.HasLaw (fun ω => idx ω t)
    (ProbabilityTheory.uniformOn (Set.univ : Set (Fin n))) P
  iIndep : ProbabilityTheory.iIndepFun (fun t ω => idx ω t) P

/-- Eqn. (eqn:lem:HU_eqn2new): every row of `H D V_k` has squared Euclidean norm at most
`2 k ln(40 n k) / n`, with `n = 2^p`. -/
def HUEqn2New (p k : ℕ) (d : Fin (2 ^ p) → ℝ) (Vk : Matrix (Fin (2 ^ p)) (Fin k) ℝ) : Prop :=
  ∀ i : Fin (2 ^ p),
    l2Norm (fun j => (normalizedHadamard p * Matrix.diagonal d * Vk) i j) ^ 2 ≤
      2 * (k : ℝ) * Real.log (40 * ((2 ^ p : ℕ) : ℝ) * (k : ℝ)) / ((2 ^ p : ℕ) : ℝ)

/-- Lemma lem:sample_lem40pfn3: assuming Eqn. (HU_eqn2new) holds for the (fixed) sign matrix
`D`, if `c ≥ 40 k ln(40 n k)/ε`, then with probability at least .95 over the uniform
sampling matrix `S`, `‖A_{k,⊥} D H S (V_kᵀ D H S)ᵀ‖_F^2 ≤ ε ‖A_{k,⊥}‖_F^2`. -/
theorem sample_lem40pfn3 {m p : ℕ} (A : Matrix (Fin m) (Fin (2 ^ p)) ℝ)
    (U : Matrix (Fin m) (Fin m) ℝ) (Sig : Matrix (Fin m) (Fin (2 ^ p)) ℝ)
    (V : Matrix (Fin (2 ^ p)) (Fin (2 ^ p)) ℝ) (s : ℕ → ℝ)
    (hSVD : IsSVD A U Sig V s)
    (k : ℕ) (hk : k ≤ 2 ^ p) (ε : ℝ) (hε : 0 < ε)
    (d : Fin (2 ^ p) → ℝ) (hd : ∀ i, d i = 1 ∨ d i = -1)
    (hHU : HUEqn2New p k d (leadingColumns k hk V))
    (c : ℕ)
    (hc : (c : ℝ) ≥ 40 * (k : ℝ) * Real.log (40 * ((2 ^ p : ℕ) : ℝ) * (k : ℝ)) / ε)
    {Ω : Type*} [MeasurableSpace Ω] (P : Measure Ω) (idx : Ω → Fin c → Fin (2 ^ p))
    (hidx : UniformSamplingIndices (2 ^ p) c P idx) :
    P {ω | frobSq (kPerp A (leadingColumns k hk V) * Matrix.diagonal d * normalizedHadamard p *
          uniformSamplingMatrix (2 ^ p) c (idx ω) *
          ((leadingColumns k hk V)ᵀ * Matrix.diagonal d * normalizedHadamard p *
            uniformSamplingMatrix (2 ^ p) c (idx ω))ᵀ) ≤
        ε * frobSq (kPerp A (leadingColumns k hk V))} ≥ (95 / 100 : ℝ≥0∞) := by
  sorry

end RNLAPhaseB.DM2017.SampleLem40pfn3