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

/-- The normalized `n × n` Hadamard matrix, `n = 2^k` (Sylvester construction
`H_{2m} = [[H_m, H_m], [H_m, -H_m]]`, scaled by `1/√n`):
`H i j = n^{-1/2} (-1)^{⟨i, j⟩}`, where `⟨i, j⟩` counts the binary digits set in both `i` and `j`. -/
noncomputable def hadamardNormalized (k : ℕ) : Matrix (Fin (2 ^ k)) (Fin (2 ^ k)) ℝ :=
  fun i j => (Real.sqrt ((2 ^ k : ℕ) : ℝ))⁻¹ *
    (-1 : ℝ) ^ ((List.range k).countP (fun b => i.val.testBit b && j.val.testBit b))

/-- Singular values of a real `r × d` matrix `M`, indexed by `Fin d`. They are the
square roots of the eigenvalues of `Mᴴ M`, zero-padded when `r < d`. Order is irrelevant
here because the lemma quantifies over all `i = 1, …, d`. -/
noncomputable def singularValues {r d : ℕ} (M : Matrix (Fin r) (Fin d) ℝ) : Fin d → ℝ :=
  fun j => Real.sqrt ((Matrix.isHermitian_conjTranspose_mul_self M).eigenvalues j)

/-- Law of the RandLeastSquares sampling indices. There are `r` i.i.d. trials with
replacement, and each trial selects an index uniformly at random from `{1, …, n}`. -/
structure UniformSamplingIndices {Ω : Type*} [MeasurableSpace Ω] (P : Measure Ω)
    (n r : ℕ) (I : Ω → Fin r → Fin n) : Prop where
  measurable : ∀ t, Measurable (fun ω => I ω t)
  uniform : ∀ t i, P {ω | I ω t = i} = ((n : ℝ≥0∞))⁻¹
  indep : ProbabilityTheory.iIndepFun (fun t ω => I ω t) P

/-- The `n × r` sampling matrix `S`. Its `t`-th column is `√(n/r) e_{I_t}`. -/
noncomputable def samplingMatrix {Ω : Type*} {n r : ℕ} (I : Ω → Fin r → Fin n) (ω : Ω) :
    Matrix (Fin n) (Fin r) ℝ :=
  fun i t => if I ω t = i then Real.sqrt ((n : ℝ) / (r : ℝ)) else 0

/-- Lemma `lem:sample_lem20pf` of Drineas–Mahoney.

The setting has `n = 2^k`, `H` the normalized Hadamard matrix and `D = diag(s)` a
(fixed, drawn) `±1` diagonal matrix. `U_A` is `n × d` with orthonormal columns.
Assume Eqn. `eqn:lem:HU_eqn2` holds for `HDU_A`, and let `r` satisfy Eqn. `eqn:rvalue`.
Then, with probability at least `.95` over the sampling matrix `S`, we have
`|1 - σ_i²(Sᵀ H D U_A)| ≤ 1 - 1/√2` for all `i = 1, …, d`. -/
theorem sample_lem20pf (k d r : ℕ) {Ω : Type*} [MeasurableSpace Ω]
    (P : Measure Ω) [IsProbabilityMeasure P]
    (U : Matrix (Fin (2 ^ k)) (Fin d) ℝ) (hU : OrthonormalColumns U)
    (s : Fin (2 ^ k) → ℝ) (hs : ∀ i, s i = 1 ∨ s i = -1)
    (hHU : ∀ i : Fin (2 ^ k),
      l2Norm ((hadamardNormalized k * Matrix.diagonal s * U) i) ^ 2 ≤
        2 * (d : ℝ) * Real.log (40 * ((2 ^ k : ℕ) : ℝ) * (d : ℝ)) / ((2 ^ k : ℕ) : ℝ))
    (hr : (48 : ℝ) ^ 2 * (d : ℝ) * Real.log (40 * ((2 ^ k : ℕ) : ℝ) * (d : ℝ)) *
        Real.log ((100 : ℝ) ^ 2 * (d : ℝ) *
          Real.log (40 * ((2 ^ k : ℕ) : ℝ) * (d : ℝ))) ≤ (r : ℝ))
    (I : Ω → Fin r → Fin (2 ^ k)) (hI : UniformSamplingIndices P (2 ^ k) r I) :
    (95 / 100 : ℝ≥0∞) ≤
      P {ω | ∀ j : Fin d,
        |1 - singularValues ((samplingMatrix I ω)ᵀ *
            (hadamardNormalized k * Matrix.diagonal s * U)) j ^ 2|
          ≤ 1 - 1 / Real.sqrt 2} := by
  sorry

end RNLAPhaseB.DM2017.SampleLem20pf