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

/-- The Frobenius norm `‖M‖_F = (∑_{i,j} M_{ij}^2)^{1/2}` of a real matrix, written
explicitly so that it does not clash with the scoped L2-operator norm instance. -/
noncomputable def frobeniusNorm {p q : ℕ} (M : Matrix (Fin p) (Fin q) ℝ) : ℝ :=
  Real.sqrt (∑ i, ∑ j, M i j ^ 2)

/-- A (full) singular value decomposition `W = U S Vᵀ` of a real `p × q` matrix:
`U`, `V` are square orthogonal, and `S` is rectangular diagonal with nonnegative
diagonal entries in nonincreasing order. -/
structure IsSingularValueDecomposition {p q : ℕ} (W : Matrix (Fin p) (Fin q) ℝ)
    (U : Matrix (Fin p) (Fin p) ℝ) (S : Matrix (Fin p) (Fin q) ℝ)
    (V : Matrix (Fin q) (Fin q) ℝ) : Prop where
  left_orthogonal : Uᵀ * U = 1 ∧ U * Uᵀ = 1
  right_orthogonal : Vᵀ * V = 1 ∧ V * Vᵀ = 1
  offdiag_zero : ∀ (i : Fin p) (j : Fin q), (i : ℕ) ≠ (j : ℕ) → S i j = 0
  diag_nonneg : ∀ (i : Fin p) (j : Fin q), (i : ℕ) = (j : ℕ) → 0 ≤ S i j
  diag_nonincreasing : ∀ (i : Fin p) (j : Fin q) (i' : Fin p) (j' : Fin q),
    (i : ℕ) = (j : ℕ) → (i' : ℕ) = (j' : ℕ) →
    (i : ℕ) ≤ (i' : ℕ) → S i' j' ≤ S i j
  factorization : W = U * S * Vᵀ

/-- The top `k` left singular vectors: the first `k` columns of the left factor `U`. -/
def leadingColumns {p k : ℕ} (hk : k ≤ p) (U : Matrix (Fin p) (Fin p) ℝ) :
    Matrix (Fin p) (Fin k) ℝ :=
  U.submatrix id (Fin.castLE hk)

/-- The rank-`k` truncation `W_k = U S_k Vᵀ` of the SVD `W = U S Vᵀ`, where `S_k`
keeps the top `k` singular values and zeroes the rest (the best rank-`k`
approximation of `W`, taken from the SAME SVD whose top-`k` left singular vectors
form `U_{W,k}`). -/
noncomputable def truncatedSVD {p q : ℕ} (k : ℕ) (U : Matrix (Fin p) (Fin p) ℝ)
    (S : Matrix (Fin p) (Fin q) ℝ) (V : Matrix (Fin q) (Fin q) ℝ) :
    Matrix (Fin p) (Fin q) ℝ :=
  U * (Matrix.of fun (i : Fin p) (j : Fin q) => if (i : ℕ) < k then S i j else 0) * Vᵀ

/-- dm2017 Lemma `lem:restate`. `U_C` (`m × ρ_C`) is an orthonormal basis of the column
span of `C`; `W = U_Cᵀ A`; `U_{W,k}` are the top `k` left singular vectors of `W`
(from an SVD `W = U S Vᵀ`); `Ũ_k = U_C U_{W,k}` is the RandLowRank output;
`(U_Cᵀ A)_k` is the rank-`k` truncation of that SVD. Conclusions:
`A - Ũ_k Ũ_kᵀ A = A - U_C (U_Cᵀ A)_k`, and
`‖A - U_C (U_Cᵀ A)_k‖_F^2 = min_{rank(Y) ≤ k} ‖A - U_C Y‖_F^2`. -/
theorem lem_restate {m n c ρ k : ℕ} (hk : k ≤ ρ)
    (A : Matrix (Fin m) (Fin n) ℝ) (C : Matrix (Fin m) (Fin c) ℝ)
    (UC : Matrix (Fin m) (Fin ρ) ℝ)
    (hUC_orth : OrthonormalColumns UC)
    (hUC_span : matrixColumnSpace UC = matrixColumnSpace C)
    (U : Matrix (Fin ρ) (Fin ρ) ℝ) (S : Matrix (Fin ρ) (Fin n) ℝ)
    (V : Matrix (Fin n) (Fin n) ℝ)
    (hSVD : IsSingularValueDecomposition (UCᵀ * A) U S V)
    (UWk : Matrix (Fin ρ) (Fin k) ℝ) (hUWk : UWk = leadingColumns hk U)
    (Utilde : Matrix (Fin m) (Fin k) ℝ) (hUtilde : Utilde = UC * UWk) :
    A - Utilde * Utildeᵀ * A = A - UC * truncatedSVD k U S V ∧
    IsLeast
      {x : ℝ | ∃ Y : Matrix (Fin ρ) (Fin n) ℝ, Y.rank ≤ k ∧
        x = frobeniusNorm (A - UC * Y) ^ 2}
      (frobeniusNorm (A - UC * truncatedSVD k U S V) ^ 2) := by
  sorry

end RNLAPhaseB.DM2017.LemRestate