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

/-- Squared Frobenius norm `‖M‖_F^2 = ∑_{i,j} M_{ij}^2` (written explicitly, since the
header scopes `‖·‖` to the operator norm). -/
def frobSq {m n : ℕ} (M : Matrix (Fin m) (Fin n) ℝ) : ℝ :=
  ∑ i, ∑ j, M i j ^ 2

/-- A thin SVD `M = U Σ Vᵀ` with `ρ` (= rank M) positive singular values in
nonincreasing order, as in Section 2.6 (`U ∈ ℝ^{m×ρ}`, `V ∈ ℝ^{n×ρ}` with orthonormal
columns). -/
structure ThinSVD {m n : ℕ} (M : Matrix (Fin m) (Fin n) ℝ) (ρ : ℕ) where
  U : Matrix (Fin m) (Fin ρ) ℝ
  V : Matrix (Fin n) (Fin ρ) ℝ
  σ : Fin ρ → ℝ
  U_orth : Uᵀ * U = 1
  V_orth : Vᵀ * V = 1
  σ_pos : ∀ i, 0 < σ i
  σ_antitone : Antitone σ
  factor : M = U * Matrix.diagonal σ * Vᵀ

/-- `A_k = U_k Σ_k V_kᵀ`: the top-`k` part of the SVD. -/
noncomputable def ThinSVD.top {m n ρ : ℕ} {M : Matrix (Fin m) (Fin n) ℝ}
    (s : ThinSVD M ρ) (k : ℕ) : Matrix (Fin m) (Fin n) ℝ :=
  s.U * Matrix.diagonal (fun i => if i.val < k then s.σ i else 0) * s.Vᵀ

/-- `A_{k,⊥} = U_{k,⊥} Σ_{k,⊥} V_{k,⊥}ᵀ`: the bottom `ρ - k` part of the SVD. -/
noncomputable def ThinSVD.bottom {m n ρ : ℕ} {M : Matrix (Fin m) (Fin n) ℝ}
    (s : ThinSVD M ρ) (k : ℕ) : Matrix (Fin m) (Fin n) ℝ :=
  s.U * Matrix.diagonal (fun i => if k ≤ i.val then s.σ i else 0) * s.Vᵀ

/-- The top `k` left singular vectors (first `k` columns of `U`), for `k ≤ ρ`. -/
def ThinSVD.topLeft {m n ρ : ℕ} {M : Matrix (Fin m) (Fin n) ℝ}
    (s : ThinSVD M ρ) (k : ℕ) (hk : k ≤ ρ) : Matrix (Fin m) (Fin k) ℝ :=
  s.U.submatrix id (Fin.castLE hk)

/-- The normalized (Sylvester/Walsh) Hadamard transform matrix of order `n = 2^p`:
`H_{ij} = (-1)^{⟨i,j⟩} / √n`, `⟨i,j⟩` the number of common 1-bits of `i` and `j`. -/
noncomputable def normalizedHadamard (n : ℕ) : Matrix (Fin n) (Fin n) ℝ :=
  fun i j => (-1 : ℝ) ^ (Nat.digits 2 (i.val &&& j.val)).sum / Real.sqrt n

/-- Step 3 of RandLowRank: `S ∈ ℝ^{n×c}` whose `t`-th column is `√(n/c) e_{i_t}`,
where `i_t = idx t` is the index selected in trial `t`. -/
noncomputable def samplingRescalingMatrix (n c : ℕ) (idx : Fin c → Fin n) :
    Matrix (Fin n) (Fin c) ℝ :=
  fun i t => if i = idx t then Real.sqrt ((n : ℝ) / c) else 0

/-- Lemma (dm2017 l_aux2). For every realization of the RandLowRank randomness
(signs `D_ii = ±1`, sampled indices `i_1,…,i_c`), with `C = A D H S`, `U_C` an
orthonormal basis of the column span of `C`, `W = U_Cᵀ A` of rank at least `k`,
`U_{W,k}` its top `k` left singular vectors and output `Ũ_k = U_C U_{W,k}`:
`‖A − Ũ_k Ũ_kᵀ A‖_F² ≤ ‖A_k − U_C U_Cᵀ A_k‖_F² + ‖A_{k,⊥}‖_F²`. -/
theorem projected_rank_k_error {m n p : ℕ} (hn : n = 2 ^ p)
    (A : Matrix (Fin m) (Fin n) ℝ) (k c : ℕ)
    (d : Fin n → ℝ) (hd : ∀ i, d i = 1 ∨ d i = -1)
    (idx : Fin c → Fin n)
    {r : ℕ} (UC : Matrix (Fin m) (Fin r) ℝ) (hUC : OrthonormalColumns UC)
    (hspan : matrixColumnSpace UC =
      matrixColumnSpace (A * Matrix.diagonal d * normalizedHadamard n *
        samplingRescalingMatrix n c idx))
    {ρW : ℕ} (svdW : ThinSVD (UCᵀ * A) ρW) (hk : k ≤ ρW)
    {ρ : ℕ} (svdA : ThinSVD A ρ) :
    frobSq (A - (UC * svdW.topLeft k hk) * (UC * svdW.topLeft k hk)ᵀ * A) ≤
      frobSq (svdA.top k - UC * UCᵀ * svdA.top k) + frobSq (svdA.bottom k) := by
  sorry

end RNLAPhaseB.DM2017.ProjectedRankKError