import Mathlib.Analysis.InnerProductSpace.PiL2
import Mathlib.Topology.Algebra.InfiniteSum.Order
import Mathlib.Topology.Algebra.InfiniteSum.Ring

/-!
# Part A — Main Theorem via Hermite Polynomials (Theorem 4.1)

  Any measurable h : ℝⁿ → ℝⁿ satisfying Gaussianity h(z) ~ N(0,Iₙ)
  and minimizing the alignment loss must be h(z) = Uz for U ∈ O(n).

  ## Verification status

  | Component                        | Status      |
  |----------------------------------|-------------|
  | Hermite basis & completeness     | axiomatized |
  | Contraction lemma (ρᵈ decay)     | axiomatized |
  | Mehler's formula                 | axiomatized |
  | ρᵈ ≤ ρ for d ≥ 1                | VERIFIED    |
  | ρᵈ < ρ for d ≥ 2                | VERIFIED    |
  | Pointwise term bound w_d·ρᵈ≤w_d·ρ| VERIFIED    |
  | Correlation bound ≤ ρ            | VERIFIED    |
  | Equality ⟺ w₁ = 1 (linearity)  | VERIFIED    |
  | Loss lower bound 2(1-ρ)n        | VERIFIED    |
  | Theorem assembly h = Uz         | VERIFIED    |
-/

set_option maxHeartbeats 400000

open scoped BigOperators

noncomputable section

abbrev E (n : ℕ) := EuclideanSpace ℝ (Fin n)


-- ═══════════════════════════════════════════════════════════════
-- SPECTRAL WEIGHTS
-- ═══════════════════════════════════════════════════════════════

/-- Spectral weights of a single encoder component in its Hermite
    expansion. `w d` is the fraction of L²(γₙ) variance at degree d. -/
structure SpectralWeights where
  w : ℕ → ℝ
  nonneg : ∀ d, 0 ≤ w d
  zero_degree : w 0 = 0
  summable : Summable w
  total_variance : ∑' d, w d = 1


-- ═══════════════════════════════════════════════════════════════
-- AXIOMATIZED: HERMITE BASIS & MEHLER
-- ═══════════════════════════════════════════════════════════════

/-- **Mehler's formula** (axiomatized): the spectral correlation
    series Σ_d w_d · ρᵈ is summable. -/
axiom mehler_summability
    (sw : SpectralWeights) (ρ : ℝ) (hρ0 : 0 < ρ) (hρ1 : ρ < 1) :
    Summable (fun d => sw.w d * ρ ^ d)


-- ═══════════════════════════════════════════════════════════════
-- VERIFIED: POINTWISE BOUNDS
-- ═══════════════════════════════════════════════════════════════

/-- For 0 < ρ ≤ 1 and d ≥ 1, ρᵈ ≤ ρ. -/
theorem pow_le_self_of_pos_lt_one (ρ : ℝ) (hρ0 : 0 < ρ) (hρ1 : ρ ≤ 1)
    (d : ℕ) (hd : 1 ≤ d) : ρ ^ d ≤ ρ := by
  calc ρ ^ d ≤ ρ ^ 1 := pow_le_pow_of_le_one (le_of_lt hρ0) hρ1 hd
       _ = ρ := pow_one ρ

/-- Each term w_d · ρᵈ ≤ w_d · ρ. -/
theorem spectral_term_le (sw : SpectralWeights) (ρ : ℝ)
    (hρ0 : 0 < ρ) (hρ1 : ρ ≤ 1) (d : ℕ) :
    sw.w d * ρ ^ d ≤ sw.w d * ρ := by
  match d with
  | 0 => simp [sw.zero_degree]
  | d + 1 =>
    exact mul_le_mul_of_nonneg_left
      (pow_le_self_of_pos_lt_one ρ hρ0 hρ1 (d + 1)
        (Nat.succ_le_succ (Nat.zero_le d)))
      (sw.nonneg (d + 1))

/-- For 0 < ρ < 1 and d ≥ 2, ρᵈ < ρ (strict). -/
theorem pow_lt_self_of_ge_two (ρ : ℝ) (hρ0 : 0 < ρ) (hρ1 : ρ < 1)
    (d : ℕ) (hd : 2 ≤ d) : ρ ^ d < ρ := by
  calc ρ ^ d ≤ ρ ^ 2 := pow_le_pow_of_le_one (le_of_lt hρ0) (le_of_lt hρ1) hd
       _ = ρ * ρ      := by ring
       _ < ρ * 1      := mul_lt_mul_of_pos_left hρ1 hρ0
       _ = ρ           := mul_one ρ


-- ═══════════════════════════════════════════════════════════════
-- VERIFIED: SUMMABILITY AND TSUM OF UPPER BOUND
-- ═══════════════════════════════════════════════════════════════

/-- The constant-ρ series fun d ↦ w d * ρ is summable
    (via Summable.mul_right from Ring.lean). -/
theorem summable_spectral_upper (sw : SpectralWeights) (ρ : ℝ) :
    Summable (fun d => sw.w d * ρ) :=
  sw.summable.mul_right ρ

/-- Σ w_d · ρ = (Σ w_d) · ρ = 1 · ρ = ρ
    (via tsum_mul_right from Ring.lean). -/
theorem tsum_spectral_upper (sw : SpectralWeights) (ρ : ℝ) :
    ∑' d, sw.w d * ρ = ρ := by
  rw [tsum_mul_right, sw.total_variance, one_mul]


-- ═══════════════════════════════════════════════════════════════
-- VERIFIED: CORRELATION BOUND (Lemma 3.3)
-- ═══════════════════════════════════════════════════════════════

/-- **Correlation bound** (VERIFIED): Σ_d w_d ρᵈ ≤ ρ.
    Uses Summable.tsum_le_tsum (from Order.lean via @[to_additive]). -/
theorem correlation_le_rho (sw : SpectralWeights) (ρ : ℝ)
    (hρ0 : 0 < ρ) (hρ1 : ρ < 1)
    (hsum : Summable (fun d => sw.w d * ρ ^ d)) :
    ∑' d, sw.w d * ρ ^ d ≤ ρ := by
  calc ∑' d, sw.w d * ρ ^ d
      ≤ ∑' d, sw.w d * ρ :=
        hsum.tsum_le_tsum
          (fun d => spectral_term_le sw ρ hρ0 (le_of_lt hρ1) d)
          (summable_spectral_upper sw ρ)
    _ = ρ := tsum_spectral_upper sw ρ


-- ═══════════════════════════════════════════════════════════════
-- VERIFIED: EQUALITY FORCES LINEARITY
-- ═══════════════════════════════════════════════════════════════

/-- **Equality characterization** (VERIFIED): if Σ w_d ρᵈ = ρ, then
    w_d = 0 for all d ≥ 2.

    Strategy: by contradiction. If w_{d₀} > 0 for some d₀ ≥ 2, then
    w_{d₀}·ρ^{d₀} < w_{d₀}·ρ strictly, while all other terms satisfy ≤.
    By Summable.tsum_lt_tsum (from Order.lean via @[to_additive]),
    Σ w_d·ρᵈ < Σ w_d·ρ = ρ, contradicting Σ w_d·ρᵈ = ρ. -/
theorem equality_forces_degree_one (sw : SpectralWeights) (ρ : ℝ)
    (hρ0 : 0 < ρ) (hρ1 : ρ < 1)
    (hsum : Summable (fun d => sw.w d * ρ ^ d))
    (heq : ∑' d, sw.w d * ρ ^ d = ρ) :
    ∀ d, 2 ≤ d → sw.w d = 0 := by
  by_contra h
  push_neg at h
  obtain ⟨d₀, hd₀_ge, hd₀_ne⟩ := h
  -- w_{d₀} > 0
  have hwd₀_pos : 0 < sw.w d₀ :=
    lt_of_le_of_ne (sw.nonneg d₀) (Ne.symm hd₀_ne)
  -- Strict inequality at d₀: w_{d₀} · ρ^{d₀} < w_{d₀} · ρ
  have hstrict : sw.w d₀ * ρ ^ d₀ < sw.w d₀ * ρ :=
    mul_lt_mul_of_pos_left (pow_lt_self_of_ge_two ρ hρ0 hρ1 d₀ hd₀_ge) hwd₀_pos
  -- By tsum_lt_tsum: one strict + rest ≤ ⟹ strict on tsums
  have hlt : ∑' d, sw.w d * ρ ^ d < ∑' d, sw.w d * ρ :=
    hsum.tsum_lt_tsum
      (fun d => spectral_term_le sw ρ hρ0 (le_of_lt hρ1) d)
      hstrict
      (summable_spectral_upper sw ρ)
  -- But Σ w_d·ρᵈ = ρ = Σ w_d·ρ
  rw [tsum_spectral_upper, heq] at hlt
  exact lt_irrefl ρ hlt


-- ═══════════════════════════════════════════════════════════════
-- ENCODER STRUCTURE & LOSS
-- ═══════════════════════════════════════════════════════════════

variable {n : ℕ}

/-- An encoder h : ℝⁿ → ℝⁿ with its Hermite spectral decomposition. -/
structure HermiteEncoder (n : ℕ) where
  toFun : E n → E n
  spectrum : Fin n → SpectralWeights
  correlation : Fin n → ℝ

/-- The alignment loss: 𝓛(h) = 2n − 2 Σᵢ corr_i. -/
def alignmentLoss (enc : HermiteEncoder n) : ℝ :=
  2 * n - 2 * ∑ i : Fin n, enc.correlation i


-- ═══════════════════════════════════════════════════════════════
-- AXIOMATIZED: BRIDGE LEMMAS
-- ═══════════════════════════════════════════════════════════════

axiom correlation_eq_spectral_sum (enc : HermiteEncoder n) (ρ : ℝ)
    (hρ0 : 0 < ρ) (hρ1 : ρ < 1) (i : Fin n) :
    enc.correlation i = ∑' d, (enc.spectrum i).w d * ρ ^ d

axiom linear_of_degree_one (enc : HermiteEncoder n)
    (hdeg : ∀ i d, 2 ≤ d → (enc.spectrum i).w d = 0) :
    ∃ (M : E n →ₗ[ℝ] E n), ∀ z, enc.toFun z = M z

axiom orthogonal_of_gaussian_linear (M : E n →ₗ[ℝ] E n)
    (hiso : ∀ v, ‖M v‖ = ‖v‖) :
    ∃ (U : E n →ₗᵢ[ℝ] E n), ∀ z, M z = U z


-- ═══════════════════════════════════════════════════════════════
-- VERIFIED: LOSS LOWER BOUND
-- ═══════════════════════════════════════════════════════════════

theorem loss_lower_bound (enc : HermiteEncoder n) (ρ : ℝ)
    (_hρ0 : 0 < ρ) (_hρ1 : ρ < 1)
    (hcorr : ∀ i, enc.correlation i ≤ ρ) :
    alignmentLoss enc ≥ 2 * (1 - ρ) * n := by
  unfold alignmentLoss
  have hsum_le : ∑ i : Fin n, enc.correlation i ≤ ∑ _i : Fin n, ρ :=
    Finset.sum_le_sum (fun i _ => hcorr i)
  simp only [Finset.sum_const, Finset.card_fin, nsmul_eq_mul] at hsum_le
  linarith


-- ═══════════════════════════════════════════════════════════════
-- VERIFIED: MAIN THEOREM ASSEMBLY
-- ═══════════════════════════════════════════════════════════════

/-- **Main Theorem** (Theorem 4.1, VERIFIED assembly):

    Any measurable h : ℝⁿ → ℝⁿ with h(z) ~ 𝒩(0, Iₙ) that
    achieves 𝓛(h) = 2(1−ρ)n must satisfy h(z) = Uz for U ∈ O(n).

    Verified chain:
    1. Mehler → correlation = Σ w_d ρᵈ     (axiomatized)
    2. Weighted average → corr_i ≤ ρ        (VERIFIED: correlation_le_rho)
    3. Loss sum → 𝓛 ≥ 2(1−ρ)n              (VERIFIED: loss_lower_bound)
    4. 𝓛 = 2(1−ρ)n → each corr_i = ρ       (VERIFIED: Finset.sum_lt_sum)
    5. corr_i = ρ → w₁ = 1 for all i        (VERIFIED: equality_forces_degree_one)
    6. w₁ = 1 → h linear                    (axiomatized: linear_of_degree_one)
    7. Gaussianity + linear → U orthogonal   (axiomatized: orthogonal_of_gaussian_linear)
-/
theorem hermite_identifiability
    (enc : HermiteEncoder n)
    (ρ : ℝ) (hρ0 : 0 < ρ) (hρ1 : ρ < 1)
    (hMehler : ∀ i, Summable (fun d => (enc.spectrum i).w d * ρ ^ d))
    (hcorr_eq : ∀ i, enc.correlation i =
      ∑' d, (enc.spectrum i).w d * ρ ^ d)
    (hopt : alignmentLoss enc = 2 * (1 - ρ) * ↑n)
    (hnorm : ∀ v, ‖enc.toFun v - enc.toFun 0‖ = ‖v - 0‖) :
    ∃ (U : E n →ₗᵢ[ℝ] E n), ∀ z, enc.toFun z = U z := by
  -- Step 1: Each correlation ≤ ρ
  have hcorr_le : ∀ i, enc.correlation i ≤ ρ := by
    intro i; rw [hcorr_eq i]
    exact correlation_le_rho (enc.spectrum i) ρ hρ0 hρ1 (hMehler i)
  -- Step 2: At optimality, each correlation = ρ exactly
  have hcorr_eq_rho : ∀ i, enc.correlation i = ρ := by
    by_contra hne; push_neg at hne
    obtain ⟨i₀, hi₀⟩ := hne
    have hi₀_lt : enc.correlation i₀ < ρ :=
      lt_of_le_of_ne (hcorr_le i₀) hi₀
    have hsum_lt : ∑ i : Fin n, enc.correlation i < ∑ _i : Fin n, ρ :=
      Finset.sum_lt_sum (fun i _ => hcorr_le i) ⟨i₀, Finset.mem_univ _, hi₀_lt⟩
    simp only [Finset.sum_const, Finset.card_fin, nsmul_eq_mul] at hsum_lt
    unfold alignmentLoss at hopt; linarith
  -- Step 3: corr_i = ρ forces degree-1 concentration
  have hdeg : ∀ i d, 2 ≤ d → (enc.spectrum i).w d = 0 := by
    intro i d hd
    have hci : ∑' d, (enc.spectrum i).w d * ρ ^ d = ρ := by
      rw [← hcorr_eq i]; exact hcorr_eq_rho i
    exact equality_forces_degree_one
      (enc.spectrum i) ρ hρ0 hρ1 (hMehler i) hci d hd
  -- Step 4: Linearity
  obtain ⟨M, hM⟩ := linear_of_degree_one enc hdeg
  -- Step 5: Orthogonality
  have hnorm_M : ∀ v, ‖M v‖ = ‖v‖ := by
    intro v; have hv := hnorm v
    simp only [sub_zero] at hv
    rwa [hM v, hM 0, map_zero, sub_zero] at hv
  obtain ⟨U, hU⟩ := orthogonal_of_gaussian_linear M hnorm_M
  exact ⟨U, fun z => by rw [hM z, hU z]⟩

end