Relative-error guarantee for RandLowRank

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

/-- Frobenius norm `‖M‖_F`, written explicitly. -/
noncomputable def frobNorm {a b : ℕ} (M : Matrix (Fin a) (Fin b) ℝ) : ℝ :=
  Real.sqrt (∑ i, ∑ j, M i j ^ 2)

/-- The normalized (Sylvester/Walsh) Hadamard matrix of order `n = 2^p`:
`H i j = (1/√n) (-1)^{⟨i,j⟩}`, where `⟨i,j⟩` counts the common one-bits of `i` and `j`. -/
noncomputable def hadamardNormalized (p : ℕ) : Matrix (Fin (2 ^ p)) (Fin (2 ^ p)) ℝ :=
  fun i j => (1 / Real.sqrt ((2 ^ p : ℕ) : ℝ)) *
    (-1 : ℝ) ^ ((Finset.range p).filter
      (fun b => (i : ℕ).testBit b = true ∧ (j : ℕ).testBit b = true)).card

/-- `V` (b × k) has orthonormal columns spanning a top-`k` right singular subspace of `A`:
it maximizes `‖A V‖_F` over all b × k matrices with orthonormal columns (Ky Fan). -/
def IsTopKRightSingularBasis {a b k : ℕ} (A : Matrix (Fin a) (Fin b) ℝ)
    (V : Matrix (Fin b) (Fin k) ℝ) : Prop :=
  OrthonormalColumns V ∧
    ∀ W : Matrix (Fin b) (Fin k) ℝ, OrthonormalColumns W → frobNorm (A * W) ≤ frobNorm (A * V)

/-- Top-`k` left singular vectors of `A` = top-`k` right singular vectors of `Aᵀ`. -/
def IsTopKLeftSingularBasis {a b k : ℕ} (A : Matrix (Fin a) (Fin b) ℝ)
    (U : Matrix (Fin a) (Fin k) ℝ) : Prop :=
  IsTopKRightSingularBasis Aᵀ U

/-- Sampling-and-rescaling matrix `S ∈ ℝ^{n×c}`: column `t` is `√(n/c) e_{idx t}`. -/
noncomputable def samplingRescalingMatrix {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

/-- Output relation of steps 7-9 of RandLowRank given `C`: `U_C` an orthonormal basis of
`range C`, `W = U_Cᵀ A` of rank at least `k`, `U_{W,k}` its top-`k` left singular vectors,
and `Ũ_k = U_C U_{W,k}`. -/
def RandLowRankOutput {m n k c : ℕ} (A : Matrix (Fin m) (Fin n) ℝ)
    (C : Matrix (Fin m) (Fin c) ℝ) (Ut : Matrix (Fin m) (Fin k) ℝ) : Prop :=
  ∃ r : ℕ, ∃ UC : Matrix (Fin m) (Fin r) ℝ,
    OrthonormalColumns UC ∧ matrixColumnSpace UC = matrixColumnSpace C ∧
    k ≤ (UCᵀ * A).rank ∧
    ∃ UWk : Matrix (Fin r) (Fin k) ℝ,
      IsTopKLeftSingularBasis (UCᵀ * A) UWk ∧ Ut = UC * UWk

/-- Eqn. (eqn:lem:HU_eqn2new): `‖(H D V_k)_{i*}‖_2^2 ≤ 2k ln(40nk)/n` for all `i`, `n = 2^p`. -/
def EqnHUeqn2new {p k : ℕ} (d : Fin (2 ^ p) → ℝ) (Vk : Matrix (Fin (2 ^ p)) (Fin k) ℝ) : Prop :=
  ∀ i : Fin (2 ^ p),
    l2Norm (fun j => (hadamardNormalized p * Matrix.diagonal d * Vk) i j) ^ 2 ≤
      2 * (k : ℝ) * Real.log (40 * ((2 ^ p : ℕ) : ℝ) * k) / ((2 ^ p : ℕ) : ℝ)

/-- Lemma lem:sample_lem40pfn4 (dm2017).

* `n = 2^p`, as required by the Hadamard transform of Section 5.1.
* `D = diag(d)` is a fixed sign diagonal. It is assumed to satisfy Eqn. (eqn:lem:HU_eqn2new) for
  the top-`k` right singular vectors `V_k` of `A`, and `A_k = A V_k V_kᵀ`.
* The probability is over the `c` i.i.d. uniform with-replacement column indices forming `S`.
* `c_0` is one constant for all inputs.

Domain conditions restored from the source:

* `k ≤ rank A`. The source's SVD notation has `V_{k,⊥} ∈ ℝ^{n×(ρ-k)}`, the bottom `ρ-k` *nonzero*
  singular vectors, so it presupposes `ρ = rank A ≥ k`.
* The positivity of `k ln n / ε²`. The threshold `c_0 (k ln n/ε²) ln(k ln n/ε²)` of
  Eqn. (cvalfinal) is only defined where the logarithm's argument is positive. Lean's junk value
  `Real.log 0 = 0` is not used to extend the source's formula to `n = 1`. -/
theorem sample_lem40pfn4 :
    ∃ c0 : ℝ, ∀ (p m k : ℕ) (A : Matrix (Fin m) (Fin (2 ^ p)) ℝ) (ε : ℝ)
      (d : Fin (2 ^ p) → ℝ) (Vk : Matrix (Fin (2 ^ p)) (Fin k) ℝ) (c : ℕ)
      (Ω : Type) [MeasurableSpace Ω] (P : Measure Ω) (idx : Ω → Fin c → Fin (2 ^ p)),
      1 ≤ k → k ≤ min m (2 ^ p) → k ≤ A.rank → 0 < ε → ε < 1 / 2 →
      (∀ i, d i = 1 ∨ d i = -1) →
      IsTopKRightSingularBasis A Vk →
      EqnHUeqn2new d Vk →
      IsProbabilityMeasure P →
      ProbabilityTheory.HasLaw idx
        (Measure.pi (fun _ : Fin c => (PMF.uniformOfFintype (Fin (2 ^ p))).toMeasure)) P →
      0 < (k : ℝ) * Real.log ((2 ^ p : ℕ) : ℝ) / ε ^ 2 →
      (c : ℝ) ≥ c0 * ((k : ℝ) * Real.log ((2 ^ p : ℕ) : ℝ) / ε ^ 2) *
        Real.log ((k : ℝ) * Real.log ((2 ^ p : ℕ) : ℝ) / ε ^ 2) →
      (85 / 100 : ℝ≥0∞) ≤ P {ω |
        (∃ Ut : Matrix (Fin m) (Fin k) ℝ, RandLowRankOutput A
            (A * Matrix.diagonal d * hadamardNormalized p * samplingRescalingMatrix (idx ω)) Ut) ∧
        ∀ Ut : Matrix (Fin m) (Fin k) ℝ, RandLowRankOutput A
            (A * Matrix.diagonal d * hadamardNormalized p * samplingRescalingMatrix (idx ω)) Ut →
          frobNorm (A - Ut * Utᵀ * A) ≤ (1 + ε) * frobNorm (A - A * Vk * Vkᵀ)} := by
  sorry

end RNLAPhaseB.DM2017.SampleLem40pfn4