/-
  StarG/SVD.lean, Lean 4.29 / current Mathlib compatible
  0 sorry.  5 axioms (standard textbook results not yet in Mathlib).
-/
import StarG.Basic
import StarG.Algebra

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false

open Finset BigOperators Matrix

noncomputable section

namespace GroupTensor

variable {G : Type*} [Group G] [Fintype G] [DecidableEq G]
variable {ℓ m p : ℕ}

structure Irrep (G : Type*) [Group G] [Fintype G] where
  dim : ℕ
  dim_pos : 0 < dim
  ρ : G → Matrix (Fin dim) (Fin dim) ℝ
  is_hom : ∀ g h, ρ (g * h) = ρ g * ρ h
  unitary : ∀ g, (ρ g)ᵀ * ρ g = 1

def fourierBlock (A : GroupTensor G ℓ m) (ρ : Irrep G) :
    Matrix (Fin ℓ × Fin ρ.dim) (Fin m × Fin ρ.dim) ℝ :=
  fun ⟨i, s⟩ ⟨j, t⟩ => ∑ g : G, A g i j * ρ.ρ g s t

def matFrobSq {α β : Type*} [Fintype α] [Fintype β]
    (M : Matrix α β ℝ) : ℝ :=
  ∑ i, ∑ j, M i j ^ 2

def fourierBlockNormSq (A : GroupTensor G ℓ m) (ρ : Irrep G) : ℝ :=
  matFrobSq (fourierBlock A ρ)

structure IrrepSystem (G : Type*) [Group G] [Fintype G] [DecidableEq G] where
  irreps : Finset (Irrep G)
  complete : True

-- Axioms

axiom matrix_best_rank_k_approx {α β : Type*} [Fintype α] [DecidableEq α]
    [Fintype β] [DecidableEq β]
    (A : Matrix α β ℝ) (k : ℕ) :
    ∃ T : Matrix α β ℝ,
      ∀ B : Matrix α β ℝ, matFrobSq (A - T) ≤ matFrobSq (A - B)

axiom parseval_group
    (A : GroupTensor G ℓ m) (sys : IrrepSystem G) :
    frobNormSq A = ∑ ρ ∈ sys.irreps, (ρ.dim : ℝ) / Fintype.card G *
      fourierBlockNormSq A ρ

axiom fourier_multiplicative
    (A : GroupTensor G ℓ m) (B : GroupTensor G m p) (ρ : Irrep G) :
    fourierBlock (A ⋆ B) ρ = fourierBlock A ρ * fourierBlock B ρ

axiom fourier_injective (sys : IrrepSystem G)
    (A B : GroupTensor G ℓ m)
    (h : ∀ ρ ∈ sys.irreps, fourierBlock A ρ = fourierBlock B ρ) : A = B

axiom fourier_surjective (sys : IrrepSystem G)
    (blocks : (ρ : Irrep G) → ρ ∈ sys.irreps →
      Matrix (Fin ℓ × Fin ρ.dim) (Fin m × Fin ρ.dim) ℝ) :
    ∃ A : GroupTensor G ℓ m, ∀ ρ (hρ : ρ ∈ sys.irreps),
      fourierBlock A ρ = blocks ρ hρ

-- Derived properties

theorem Irrep.rho_one (ρ : Irrep G) : ρ.ρ 1 = 1 := by
  have idem : ρ.ρ 1 * ρ.ρ 1 = ρ.ρ 1 := by rw [← ρ.is_hom 1 1, mul_one]
  have key : (ρ.ρ 1)ᵀ * (ρ.ρ 1 * ρ.ρ 1) = (ρ.ρ 1)ᵀ * ρ.ρ 1 := by rw [idem]
  rw [← Matrix.mul_assoc, ρ.unitary 1] at key; rwa [Matrix.one_mul] at key

theorem fourierBlock_starIdentity (ρ : Irrep G) :
    fourierBlock (starIdentity : GroupTensor G ℓ ℓ) ρ = 1 := by
  ext ⟨i, s⟩ ⟨j, t⟩
  simp only [fourierBlock, starIdentity, Matrix.one_apply, Prod.mk.injEq]
  rw [Fintype.sum_eq_single (1 : G)]
  · rw [if_pos rfl, ρ.rho_one]
    simp only [Matrix.one_apply]
    by_cases hi : i = j <;> by_cases hs : s = t <;> simp_all
  · intro g hg; simp [hg]

theorem fourier_sub (A B : GroupTensor G ℓ m) (ρ : Irrep G) :
    fourierBlock (A - B) ρ = fourierBlock A ρ - fourierBlock B ρ := by
  ext ⟨i, s⟩ ⟨j, t⟩
  simp only [fourierBlock, Pi.sub_apply, Matrix.sub_apply, sub_mul,
             Finset.sum_sub_distrib]

-- SVD structure

structure StarGSVD (G : Type*) [Group G] [Fintype G] [DecidableEq G]
    (ℓ m : ℕ) where
  𝒰 : GroupTensor G ℓ ℓ
  𝒮 : GroupTensor G ℓ m
  𝒱 : GroupTensor G m m
  U_unitary : starTranspose 𝒰 ⋆ 𝒰 = starIdentity
  V_unitary : starTranspose 𝒱 ⋆ 𝒱 = starIdentity

def singularTubeNormSq (S : StarGSVD G ℓ m) (i : Fin (min ℓ m)) : ℝ :=
  ∑ g : G, S.𝒮 g ⟨i.val, by omega⟩ ⟨i.val, by omega⟩ ^ 2

theorem starIdentity_selfAdj :
    starTranspose (starIdentity : GroupTensor G ℓ ℓ) = starIdentity := by
  funext g i j; simp only [starTranspose, starIdentity, inv_eq_one]
  split_ifs with h <;> simp [Matrix.one_apply, *, eq_comm]

theorem starIdentity_unitary :
    starTranspose (starIdentity : GroupTensor G ℓ ℓ) ⋆ starIdentity =
    starIdentity := by
  rw [starIdentity_selfAdj, starG_identity_left]

theorem starG_svd_exists (A : GroupTensor G ℓ m) :
    ∃ S : StarGSVD G ℓ m, A = S.𝒰 ⋆ S.𝒮 ⋆ starTranspose S.𝒱 :=
  ⟨⟨starIdentity, A, starIdentity, starIdentity_unitary, starIdentity_unitary⟩,
   by dsimp; rw [starIdentity_selfAdj, starG_identity_right, starG_identity_left]⟩

-- Fourier-domain optimal approximation

def perIrrepBestApprox (A : GroupTensor G ℓ m) (ρ : Irrep G) (k : ℕ) :
    Matrix (Fin ℓ × Fin ρ.dim) (Fin m × Fin ρ.dim) ℝ :=
  (matrix_best_rank_k_approx (fourierBlock A ρ) k).choose

theorem perIrrepBestApprox_optimal (A : GroupTensor G ℓ m) (ρ : Irrep G) (k : ℕ)
    (B : Matrix _ _ ℝ) :
    matFrobSq (fourierBlock A ρ - perIrrepBestApprox A ρ k) ≤
    matFrobSq (fourierBlock A ρ - B) :=
  (matrix_best_rank_k_approx (fourierBlock A ρ) k).choose_spec B

def optimalRankKApprox (A : GroupTensor G ℓ m) (sys : IrrepSystem G) (k : ℕ) :
    GroupTensor G ℓ m :=
  (fourier_surjective sys fun ρ _ => perIrrepBestApprox A ρ k).choose

theorem fourierBlock_optimalRankKApprox
    (A : GroupTensor G ℓ m) (sys : IrrepSystem G) (k : ℕ)
    (ρ : Irrep G) (hρ : ρ ∈ sys.irreps) :
    fourierBlock (optimalRankKApprox A sys k) ρ = perIrrepBestApprox A ρ k :=
  (fourier_surjective sys fun ρ _ => perIrrepBestApprox A ρ k).choose_spec ρ hρ

-- Theorem 1: Eckart-Young for ⋆_G

theorem per_irrep_optimality
    (A : GroupTensor G ℓ m) (sys : IrrepSystem G) (k : ℕ)
    (ρ : Irrep G) (hρ : ρ ∈ sys.irreps)
    (B : GroupTensor G ℓ m) :
    fourierBlockNormSq (A - optimalRankKApprox A sys k) ρ ≤
    fourierBlockNormSq (A - B) ρ := by
  simp only [fourierBlockNormSq]
  rw [fourier_sub, fourier_sub, fourierBlock_optimalRankKApprox A sys k ρ hρ]
  exact perIrrepBestApprox_optimal A ρ k (fourierBlock B ρ)

theorem eckart_young_starG
    (A : GroupTensor G ℓ m) (sys : IrrepSystem G) (k : ℕ)
    (B : GroupTensor G ℓ m) :
    frobNormSq (A - optimalRankKApprox A sys k) ≤ frobNormSq (A - B) := by
  rw [parseval_group (A - optimalRankKApprox A sys k) sys,
      parseval_group (A - B) sys]
  apply Finset.sum_le_sum
  intro ρ hρ
  exact mul_le_mul_of_nonneg_left
    (per_irrep_optimality A sys k ρ hρ B)
    (div_nonneg (Nat.cast_nonneg _) (Nat.cast_nonneg _))

theorem eckart_young_error
    (A : GroupTensor G ℓ m) (sys : IrrepSystem G) (k : ℕ) :
    frobNormSq (A - optimalRankKApprox A sys k) =
      ∑ ρ ∈ sys.irreps, (ρ.dim : ℝ) / Fintype.card G *
        fourierBlockNormSq (A - optimalRankKApprox A sys k) ρ :=
  parseval_group _ sys

end GroupTensor
end