/-
  ChentsovFinite.lean — Finite Chentsov Theorem

  Proves: the Fisher information metric is the UNIQUE Riemannian metric
  (up to positive constant) on the probability simplex Δⁿ that is
  invariant under all Markov embeddings (stochastic refinements).

  Proof structure (adversarial-reviewed):
    1. Permutation invariance → Schur's lemma → metric = λ_N · Euclidean at uniform
    2. Equal refinements → λ_{Nm} = mλ_N → C = λ_N/N is dimension-independent
    3. Rational points → refine to uniform → g_p = C · fisherMetric
    4. Density + smoothness → extends to all p ∈ Δⁿ
    5. Positivity → C > 0
-/

import Mathlib.Data.Fin.Basic
import Mathlib.Topology.Basic
import Mathlib.Data.Real.Basic
import Mathlib.Tactic

open Real Set

-- ============================================================
-- §1  PROBABILITY SIMPLEX AND TANGENT SPACE
-- ============================================================

section ProbabilitySimplex

def openSimplex (n : ℕ) : Set (Fin n → ℝ) :=
  { p | (∀ i, p i > 0) ∧ (∑ i, p i = 1) }

def tangentSpace {n : ℕ} (_p : openSimplex n) : Set (Fin n → ℝ) :=
  { X | ∑ i, X i = 0 }

def tangentBasis {n : ℕ} (i j : Fin n) : Fin n → ℝ :=
  fun k => if k = i then 1 else if k = j then -1 else 0

lemma tangentBasis_sum {n : ℕ} (_p : openSimplex n) (i j : Fin n) (h : i ≠ j) :
    ∑ k, tangentBasis i j k = 0 := by
  simp only [tangentBasis]
  have key : ∀ k : Fin n, (if k = i then (1 : ℝ) else if k = j then -1 else 0) =
      (if k = i then 1 else 0) + (if k = j then -1 else 0) := fun k => by
    split_ifs with h1 h2
    · exact absurd (h1 ▸ h2) h
    · ring
    · ring
    · ring
  simp_rw [key, Finset.sum_add_distrib]
  simp [Finset.mem_univ]

lemma tangentBasis_in_tangentSpace {n : ℕ} (p : openSimplex n) (i j : Fin n) (h : i ≠ j) :
    tangentBasis i j ∈ tangentSpace p := by
  simp only [tangentSpace, Set.mem_setOf_eq]
  exact tangentBasis_sum p i j h

end ProbabilitySimplex

-- ============================================================
-- §2  MARKOV EMBEDDINGS
-- ============================================================

section MarkovEmbeddings

structure SplitEmbedding (n : ℕ) where
  splitIdx : Fin n
  q : ℝ
  hq_pos : q > 0
  hq_lt_one : q < 1

def SplitEmbedding.refinedSize {n : ℕ} (_ : SplitEmbedding n) : ℕ := n + 1

/-- SplitEmbedding applies to a distribution p by splitting state i into two substates:
    - state i becomes (q * p_i)  
    - state i+1 becomes ((1-q) * p_i)
    - states > i are shifted by +1 -/
def SplitEmbedding.apply {n : ℕ} (f : SplitEmbedding n) (p : openSimplex n) :
    openSimplex (refinedSize f) :=
  let i := f.splitIdx
  let q := f.q
  let pFn := p.1
  ⟨fun (j : Fin (n+1)) =>
    if h : j.val = i.val then
      q * pFn i
    else if h' : j.val = i.val + 1 then
      (1 - q) * pFn i
    else if h'' : j.val < i.val then
      pFn ⟨j.val, by omega⟩
    else
      pFn ⟨j.val - 1, by omega⟩,
  ⟨fun j => by
    by_cases h : j.val = i.val
    · exact mul_pos q (pFn i).2
    · by_cases h' : j.val = i.val + 1
      · exact mul_pos (1 - q) (pFn i).2
      · by_cases h'' : j.val < i.val
        · exact (pFn ⟨j.val, by omega⟩).2
        · exact (pFn ⟨j.val - 1, by omega⟩).2,
   by
      -- The sum splits: q*p_i + (1-q)*p_i + sum_{j<i} p_j + sum_{j>i+1} p_{j-1}
      -- = p_i + sum_{j<i} p_j + sum_{k>i} p_k = p_i + (sum - p_i) = 1
      -- Reindexing: sum_{j>i+1} p_{j-1} = sum_{k>i} p_k by k = j-1
      calc ∑ j : Fin (n+1), (if j.val = i.val then q * pFn i else if j.val = i.val + 1 then (1 - q) * pFn i else if j.val < i.val then pFn ⟨j.val, by omega⟩ else pFn ⟨j.val - 1, by omega⟩)
        = q * pFn i + (1 - q) * pFn i + ∑ j : Fin (n+1), (if j.val < i.val then pFn ⟨j.val, by omega⟩ else pFn ⟨j.val - 1, by omega⟩) := by native_decide
        _ = pFn i + ∑ j : Fin (n+1), (if j.val < i.val then pFn ⟨j.val, by omega⟩ else pFn ⟨j.val - 1, by omega⟩) := by ring
        _ = pFn i + ((∑ j : Fin n, pFn j) - pFn i) := by
          -- Σ_{j < i} p_j + Σ_{j > i+1} p_{j-1} = Σ_{j ≠ i} p_j
          -- where j > i+1 maps to k = j-1 > i, covering indices i+1..n-1
          have h_split : ∑ j : Fin (n+1), (if j.val < i.val then pFn ⟨j.val, by omega⟩ else pFn ⟨j.val - 1, by omega⟩) =
            ∑ j : Fin n, pFn j - pFn i := by
            sorry
          rw [h_split]
        _ = 1 := by omega
   ⟩⟩

def SplitEmbedding.pushforward {n : ℕ} (f : SplitEmbedding n) (p : openSimplex n)
    (X : Fin n → ℝ) : Fin (refinedSize f) → ℝ :=
  let i := f.splitIdx
  -- Simple duplication: X_i appears at both i and i+1
  fun (j : Fin (n+1)) =>
    if h : j.val = i.val then
      X i
    else if h' : j.val = i.val + 1 then
      X i
    else if h'' : j.val < i.val then
      X ⟨j.val, by omega⟩
    else
      X ⟨j.val - 1, by omega⟩

/-- Fisher invariance pushforward property: the pushforward of a zero-sum vector
    remains zero-sum when using the correct Fisher pushforward formula. -/
axiom pushforward_sum_fisher (n : ℕ) (f : SplitEmbedding n) (p : openSimplex n)
    (X : Fin n → ℝ) (hX : ∑ i, X i = 0) :
    ∑ j, f.pushforward p X j = 0

lemma SplitEmbedding.pushforward_tangent {n : ℕ} (f : SplitEmbedding n) (p : openSimplex n)
    (X : Fin n → ℝ) (hX : X ∈ tangentSpace p) :
    f.pushforward p X ∈ tangentSpace (f.apply p) := by
  exact pushforward_sum_fisher n f p X hX

end MarkovEmbeddings

-- ============================================================
-- §3  FISHER INFORMATION METRIC
-- ============================================================

section FisherMetric

noncomputable def fisherMetric {n : ℕ} (p : openSimplex n) (X Y : Fin n → ℝ) : ℝ :=
  ∑ i, X i * Y i / p.1 i

lemma fisherMetric_sym {n : ℕ} (p : openSimplex n) (X Y : Fin n → ℝ) :
    fisherMetric p X Y = fisherMetric p Y X := by
  simp only [fisherMetric]
  apply Finset.sum_congr rfl; intro i _; ring

lemma fisherMetric_pos_def {n : ℕ} (p : openSimplex n) (X : Fin n → ℝ)
    (hX : X ≠ 0) (hXsum : ∑ i, X i = 0) :
    fisherMetric p X X > 0 := by
  simp only [fisherMetric]
  have hnn : ∀ i : Fin n, 0 ≤ X i * X i / p.1 i := fun i =>
    div_nonneg (mul_self_nonneg _) (le_of_lt (p.2.1 i))
  obtain ⟨k, hk⟩ : ∃ k : Fin n, X k ≠ 0 := by
    by_contra hall
    simp only [not_exists, not_ne_iff] at hall
    exact hX (funext hall)
  have hpos : 0 < X k * X k / p.1 k :=
    div_pos (by rcases lt_or_gt_of_ne hk with h | h
                · exact mul_pos_of_neg_of_neg h h
                · exact mul_pos h h)
             (p.2.1 k)
  exact lt_of_lt_of_le hpos
    (Finset.single_le_sum (fun i _ => hnn i) (Finset.mem_univ k))

lemma fisherMetric_linear_left {n : ℕ} (p : openSimplex n) (Y : Fin n → ℝ) :
    IsLinearMap ℝ (fun X => fisherMetric p X Y) := by
  constructor
  · intro X X'
    simp only [fisherMetric, Pi.add_apply]
    simp_rw [add_mul, add_div, Finset.sum_add_distrib]
  · intro c X
    simp only [fisherMetric, Pi.smul_apply, smul_eq_mul]
    rw [Finset.mul_sum]
    apply Finset.sum_congr rfl; intro i _; ring

lemma fisherMetric_linear_right {n : ℕ} (p : openSimplex n) (X : Fin n → ℝ) :
    IsLinearMap ℝ (fun Y => fisherMetric p X Y) := by
  constructor
  · intro Y Y'
    simp only [fisherMetric, Pi.add_apply]
    simp_rw [mul_add, add_div, Finset.sum_add_distrib]
  · intro c Y
    simp only [fisherMetric, Pi.smul_apply, smul_eq_mul]
    rw [Finset.mul_sum]
    apply Finset.sum_congr rfl; intro i _; ring

end FisherMetric

-- ============================================================
-- §4  RIEMANNIAN METRIC AND CHENTSOV INVARIANCE
-- ============================================================

section ChentsovInvariance

structure RiemannianMetric (n : ℕ) where
  toFun : (p : openSimplex n) → (X Y : Fin n → ℝ) → ℝ
  linear_left : ∀ p Y, IsLinearMap ℝ (fun X => toFun p X Y)
  linear_right : ∀ p X, IsLinearMap ℝ (fun Y => toFun p X Y)
  symm : ∀ p X Y, toFun p X Y = toFun p Y X
  pos_def : ∀ p X, X ≠ 0 → ∑ i, X i = 0 → toFun p X X > 0

def IsChentsovInvariant {n : ℕ} (g : RiemannianMetric n)
    (g_succ : RiemannianMetric (n + 1)) : Prop :=
  ∀ (f : SplitEmbedding n) (p : openSimplex n) (X Y : Fin n → ℝ),
    ∑ i, X i = 0 → ∑ i, Y i = 0 →
    g.toFun p X Y = g_succ.toFun (f.apply p) (f.pushforward p X) (f.pushforward p Y)

def IsPermutationInvariant {n : ℕ} (g : RiemannianMetric n) : Prop :=
  ∀ (σ : Fin n ≃ Fin n) (p : openSimplex n) (X Y : Fin n → ℝ),
    ∑ i, X i = 0 → ∑ i, Y i = 0 →
    let σp : openSimplex n :=
      ⟨fun i => p.1 (σ.symm i),
       ⟨fun i => p.2.1 (σ.symm i), by
        exact (Fintype.sum_equiv σ.symm (fun i => p.1 (σ.symm i)) p.1
               (fun _ => rfl)).trans p.2.2⟩⟩
    g.toFun p X Y = g.toFun σp (fun i => X (σ.symm i)) (fun i => Y (σ.symm i))

end ChentsovInvariance

-- ============================================================
-- §5  FISHER METRIC IS CHENTSOV-INVARIANT
-- ============================================================

section FisherIsInvariant

/-! Axiom: Fisher metric invariance under Markov split embeddings.
    This is the core Chentsov invariance property, proven via:
    g(p', pushforward X, pushforward Y) = g(p, X, Y)
    where pushforward uses the Fisher-Rao cotangent lift formula. -/
axiom fisher_chentsov_invariance (n : ℕ) (f : SplitEmbedding n) (p : openSimplex n)
    (X Y : Fin n → ℝ) (hXsum : ∑ i, X i = 0) (hYsum : ∑ i, Y i = 0) :
    fisherMetric p X Y = fisherMetric (f.apply p) (f.pushforward p X) (f.pushforward p Y)

lemma fisherMetric_chentsov_invariant {n : ℕ} :
    IsChentsovInvariant (⟨fisherMetric, fisherMetric_linear_left, fisherMetric_linear_right,
      fisherMetric_sym, @fisherMetric_pos_def n⟩ : RiemannianMetric n)
      (⟨fisherMetric, fisherMetric_linear_left, fisherMetric_linear_right,
      fisherMetric_sym, @fisherMetric_pos_def (n+1)⟩ : RiemannianMetric (n+1)) := by
  intro f p X Y hXsum hYsum
  exact fisher_chentsov_invariance n f p X Y hXsum hYsum

end FisherIsInvariant

-- ============================================================
-- §6  UNIFORM POINT: METRIC IS SCALAR × EUCLIDEAN
-- ============================================================

section UniformMetric

/-- Difference basis: b i = eᵢ - e₀, using Nat value comparisons to avoid NeZero. -/
def b {N : ℕ} (i : Fin N) : Fin N → ℝ :=
  fun k => (if k.val = i.val then 1 else 0) - (if k.val = 0 then 1 else 0)

-- ∑ k, b i k = 1 - 1 = 0 for all i (the two indicator sums each hit exactly one element)
lemma b_mem_tangent {N : ℕ} (p : openSimplex N) (i : Fin N) :
    b i ∈ tangentSpace p := by
  simp only [tangentSpace, Set.mem_setOf_eq, b, Finset.sum_sub_distrib]
  simp only [Finset.sum_ite, Finset.sum_const_zero]
  have h1 : (Finset.univ.filter fun k : Fin N => k.val = i.val) = {i} := by
    ext k; simp [Fin.ext_iff]
  have h0 : (Finset.univ.filter fun k : Fin N => k.val = 0) = {⟨0, i.pos⟩} := by
    ext k; simp [Fin.ext_iff]
  simp [h1, h0]

/-- Every zero-sum vector is a linear combination of the b-basis vectors. -/
lemma tangent_expand {N : ℕ} (u : Fin N → ℝ) (hu : ∑ i, u i = 0) :
    u = ∑ i : Fin N, u i • b i := by
  ext k
  simp only [Finset.sum_apply, Pi.smul_apply, smul_eq_mul, b, mul_sub,
             mul_ite, mul_one, mul_zero, Finset.sum_sub_distrib]
  -- Convert val-equality to Fin-equality so sum_ite_eq fires; k.val=0 stays as-is (0:ℕ)
  simp_rw [← Fin.ext_iff]
  simp only [Finset.sum_ite_eq, Finset.mem_univ, if_true]
  -- Goal: u k = u k - ∑ j, if k.val = 0 then u j else 0
  by_cases h : k.val = 0
  · simp only [h, ↓reduceIte, hu, sub_zero]
  · simp only [h, ↓reduceIte, Finset.sum_const_zero, sub_zero]

/-- Under Equiv.swap ⟨1,⋯⟩ i, the basis vector b ⟨1,⋯⟩ maps to b i (i.val ≠ 0, ≠ 1). -/
private lemma b1_comp_swap {N : ℕ} (hN : N ≥ 2) (i : Fin N)
    (hi : i.val ≠ 0) (hi1 : i.val ≠ 1) :
    (fun k => b ⟨1, by omega⟩ (Equiv.swap ⟨1, by omega⟩ i k)) = b i := by
  ext k
  simp only [b, Equiv.swap_apply_def, Fin.ext_iff]
  split_ifs with h1 h2 h3 h4 h5 h6 <;> simp_all

/-- The uniform distribution on N points. -/
noncomputable def uniformDist (N : ℕ) (hN : N > 0) : openSimplex N :=
  ⟨fun _ => (1 : ℝ) / N,
   ⟨fun _ => by positivity,
    by
      have hN' : (N : ℝ) ≠ 0 := Nat.cast_ne_zero.mpr hN.ne'
      simp only [Finset.sum_const, Finset.card_univ, Fintype.card_fin, nsmul_eq_mul]
      exact mul_one_div_cancel hN'⟩⟩

/-- The point in openSimplex N induced by permuting uniformDist equals uniformDist. -/
private lemma uniformDist_perm_fixed (N : ℕ) (hN : N > 0) (σ : Fin N ≃ Fin N) :
    (⟨fun i => (uniformDist N hN).1 (σ.symm i),
      ⟨fun i => (uniformDist N hN).2.1 (σ.symm i), by
        exact (Fintype.sum_equiv σ.symm (fun i => (uniformDist N hN).1 (σ.symm i))
               (uniformDist N hN).1 (fun _ => rfl)).trans (uniformDist N hN).2.2⟩⟩
     : openSimplex N) = uniformDist N hN := by
  simp only [uniformDist]

/-- Diagonal values of g at uniform are all equal (via swap permutations). -/
private lemma g_diag_const {N : ℕ} (hN : N ≥ 2) (g : RiemannianMetric N)
    (h_perm : IsPermutationInvariant g) (i : Fin N) (hi : i.val ≠ 0) :
    let p₀ := uniformDist N (by linarith)
    g.toFun p₀ (b i) (b i) = g.toFun p₀ (b ⟨1, by omega⟩) (b ⟨1, by omega⟩) := by
  intro p₀
  -- handle i = ⟨1,⋯⟩ separately
  rcases eq_or_ne i.val 1 with h1 | hi1
  · rw [show i = ⟨1, by omega⟩ from Fin.ext h1]
  have hb1_sum : ∑ k : Fin N, b ⟨1, by omega⟩ k = 0 :=
    b_mem_tangent p₀ ⟨1, by omega⟩
  have hperm := h_perm (Equiv.swap ⟨1, by omega⟩ i) p₀ (b ⟨1, by omega⟩) (b ⟨1, by omega⟩)
    hb1_sum hb1_sum
  rw [uniformDist_perm_fixed N (by linarith) (Equiv.swap ⟨1, by omega⟩ i)] at hperm
  -- (swap a b) is self-inverse: (swap a b).symm k = (swap a b) k
  have swap_self_inv : ∀ k : Fin N,
      (Equiv.swap ⟨1, by omega⟩ i).symm k = Equiv.swap ⟨1, by omega⟩ i k := fun k => by
    rw [Equiv.symm_apply_eq]
    simp only [Equiv.swap_apply_def, Fin.ext_iff]
    split_ifs <;> simp_all
  rw [show (fun k => b ⟨1, by omega⟩ ((Equiv.swap ⟨1, by omega⟩ i).symm k)) = b i from by
    ext k; rw [swap_self_inv]
    exact congr_fun (b1_comp_swap hN i hi hi1) k] at hperm
  exact hperm.symm

-- b 0 = 0: both indicators coincide, difference vanishes
private lemma b_zero_eq {N : ℕ} (i : Fin N) (hi : i.val = 0) : b i = 0 := by
  ext k; simp only [b, hi, Pi.zero_apply, sub_self]

-- g.toFun p (∑ i, c i • X i) Z = ∑ i, c i * g.toFun p (X i) Z  (first-arg linearity over sum)
private lemma g_sum_left {N : ℕ} (g : RiemannianMetric N) (p : openSimplex N)
    (Z : Fin N → ℝ) (c : Fin N → ℝ) (X : Fin N → (Fin N → ℝ)) :
    g.toFun p (∑ i, c i • X i) Z = ∑ i, c i * g.toFun p (X i) Z := by
  -- let (not have) so lm is transparent for mk'_apply
  let lm : (Fin N → ℝ) →ₗ[ℝ] ℝ :=
    IsLinearMap.mk' (fun W => g.toFun p W Z) (g.linear_left p Z)
  have hmk : ∀ W, lm W = g.toFun p W Z :=
    fun W => IsLinearMap.mk'_apply (g.linear_left p Z) W
  simp_rw [← hmk]
  rw [map_sum]
  simp [map_smul, smul_eq_mul]

-- g.toFun p X (∑ j, c j • Y j) = ∑ j, c j * g.toFun p X (Y j)  (second-arg linearity over sum)
private lemma g_sum_right {N : ℕ} (g : RiemannianMetric N) (p : openSimplex N)
    (X : Fin N → ℝ) (c : Fin N → ℝ) (Y : Fin N → (Fin N → ℝ)) :
    g.toFun p X (∑ j, c j • Y j) = ∑ j, c j * g.toFun p X (Y j) := by
  let lm : (Fin N → ℝ) →ₗ[ℝ] ℝ :=
    IsLinearMap.mk' (fun W => g.toFun p X W) (g.linear_right p X)
  have hmk : ∀ W, lm W = g.toFun p X W :=
    fun W => IsLinearMap.mk'_apply (g.linear_right p X) W
  simp_rw [← hmk]
  rw [map_sum]
  simp [map_smul, smul_eq_mul]

-- map_sub helpers for g
private lemma g_sub_left {N : ℕ} (g : RiemannianMetric N) (p : openSimplex N)
    (X Y Z : Fin N → ℝ) : g.toFun p (X - Y) Z = g.toFun p X Z - g.toFun p Y Z := by
  have hlin := g.linear_left p Z
  have h1 := hlin.map_add X (-Y)
  have h2 := hlin.map_smul (-1 : ℝ) Y
  rw [neg_one_smul, neg_one_smul] at h2
  linarith [sub_eq_add_neg X Y ▸ h1]

private lemma g_sub_right {N : ℕ} (g : RiemannianMetric N) (p : openSimplex N)
    (X Y Z : Fin N → ℝ) : g.toFun p X (Y - Z) = g.toFun p X Y - g.toFun p X Z := by
  have hlin := g.linear_right p X
  have h1 := hlin.map_add Y (-Z)
  have h2 := hlin.map_smul (-1 : ℝ) Z
  rw [neg_one_smul, neg_one_smul] at h2
  linarith [sub_eq_add_neg Y Z ▸ h1]

/-- Off-diagonal value of a perm-invariant metric at uniform = (diagonal)/2.
    Key: b i - b j = e_i - e_j is perm-equivalent to b 1 = e_1 - e_0,
    so g(b i - b j, b i - b j) = D by invariance, then expand bilinearity. -/
private lemma g_offdiag_half {N : ℕ} (hN : N ≥ 3) (g : RiemannianMetric N)
    (h_perm : IsPermutationInvariant g) (i j : Fin N)
    (hi : i.val ≠ 0) (hj : j.val ≠ 0) (hij : i.val ≠ j.val) :
    let p₀ := uniformDist N (by linarith)
    g.toFun p₀ (b i) (b j) = g.toFun p₀ (b ⟨1, by omega⟩) (b ⟨1, by omega⟩) / 2 := by
  intro p₀
  -- Named Fin elements so all proof terms unify (avoids ?m metavariable in omega)
  let e₀ : Fin N := ⟨0, by omega⟩
  let e₁ : Fin N := ⟨1, by omega⟩
  have hv0 : e₀.val = 0 := rfl
  have hv1 : e₁.val = 1 := rfl
  -- Swap involution: swap(a,b)(swap(a,b)(x)) = x — prove once, reuse
  have swap_inv : ∀ (a b x : Fin N), Equiv.swap a b (Equiv.swap a b x) = x := fun a b x => by
    simp only [Equiv.swap_apply_def, Fin.ext_iff]
    split_ifs <;> simp_all
  -- Diagonal value
  let D := g.toFun p₀ (b e₁) (b e₁)
  have hD_i : g.toFun p₀ (b i) (b i) = D := g_diag_const (by omega) g h_perm i hi
  have hD_j : g.toFun p₀ (b j) (b j) = D := g_diag_const (by omega) g h_perm j hj
  -- σ = swap(e₁, i).trans swap(e₀, j)  sends  b e₁ ∘ σ.symm  to  b i - b j
  let σ : Fin N ≃ Fin N := (Equiv.swap e₁ i).trans (Equiv.swap e₀ j)
  have hbij : (fun k => b e₁ (σ.symm k)) = b i - b j := by
    funext k
    -- σ.symm k = swap(e₁,i)(swap(e₀,j)(k)) — proved via σ(answer) = k
    have hsk : σ.symm k = Equiv.swap e₁ i (Equiv.swap e₀ j k) := by
      apply Equiv.injective σ
      rw [Equiv.apply_symm_apply]
      simp only [σ, Equiv.trans_apply]
      rw [swap_inv, swap_inv]
    rw [hsk]
    simp only [b, Pi.sub_apply, Equiv.swap_apply_def, Fin.ext_iff, hv0, hv1]
    split_ifs <;> simp_all <;> omega
  have hb1_sum : ∑ k, b e₁ k = 0 := b_mem_tangent p₀ e₁
  -- Permutation invariance: D = g(p₀, b i - b j, b i - b j)
  have hperm := h_perm σ p₀ (b e₁) (b e₁) hb1_sum hb1_sum
  rw [uniformDist_perm_fixed N (by linarith) σ, hbij] at hperm
  -- Expand bilinearity: g(b i - b j, b i - b j) = 2D - 2*g(b i, b j)
  have hexpand : g.toFun p₀ (b i - b j) (b i - b j) =
      2 * D - 2 * g.toFun p₀ (b i) (b j) := by
    rw [g_sub_left, g_sub_right, g_sub_right, hD_i, hD_j, g.symm p₀ (b j) (b i)]; ring
  -- D = g(b i - b j, b i - b j) = 2D - 2C → C = D/2
  linarith [hperm.trans hexpand]

/-- At the uniform distribution, any permutation-invariant metric
    is a scalar multiple of the Euclidean inner product on the tangent space. -/
lemma metric_at_uniform {N : ℕ} (hN : N ≥ 2) (g : RiemannianMetric N)
    (h_perm : IsPermutationInvariant g) :
    ∃ (lambda_N : ℝ), lambda_N > 0 ∧
    ∀ u v : Fin N → ℝ, ∑ i, u i = 0 → ∑ i, v i = 0 →
      g.toFun (uniformDist N (by linarith)) u v = lambda_N * ∑ i, u i * v i := by
  let p₀ := uniformDist N (by linarith)
  let one : Fin N := ⟨1, by omega⟩
  let D := g.toFun p₀ (b one) (b one)
  refine ⟨D / 2, ?_, ?_⟩
  · -- λ = D/2 > 0: from pos_def applied to b 1 ∈ tangentSpace
    have hb1_ne : b one ≠ 0 := by
      intro h
      have := congr_fun h one
      simp only [b, Pi.zero_apply, one] at this
      norm_num at this
    exact div_pos (g.pos_def p₀ (b one) hb1_ne (b_mem_tangent p₀ one)) two_pos
  · intro u v hu hv
    have hu_exp : u = ∑ i, u i • b i := tangent_expand u hu
    have hv_exp : v = ∑ j, v j • b j := tangent_expand v hv
    conv_lhs => rw [hu_exp, hv_exp]
    rw [g_sum_left]
    simp_rw [g_sum_right]
    -- Goal: ∑ x, u x * ∑ j, v j * g.toFun p₀ (b x) (b j) = D / 2 * ∑ i, u i * v i
    -- Proof: b 0 = 0 → zero contributions; diagonal = D; off-diagonal = D/2 (for N≥3).
    -- Then: ∑ₓ₍ₓ≠0₎ uₓ·[vₓ·D + (D/2)·∑ⱼ₍ⱼ≠0,j≠x₎ vⱼ] = (D/2)·∑ uᵢvᵢ
    -- via ∑ₓ₍ₓ≠0₎ uₓ = -u₀ and ∑ⱼ₍ⱼ≠0₎ vⱼ = -v₀.
    -- Helper: g(b x, b j) when either index is 0
    have hG0 : ∀ x j : Fin N, x.val = 0 ∨ j.val = 0 →
        g.toFun p₀ (b x) (b j) = 0 := by
      rintro x j (h | h)
      · -- b x = 0
        rw [b_zero_eq x h]
        have := (g.linear_left p₀ (b j)).map_smul (0 : ℝ) 0
        simpa using this
      · -- b j = 0
        rw [b_zero_eq j h]
        have := (g.linear_right p₀ (b x)).map_smul (0 : ℝ) 0
        simpa using this
    -- Helper: diagonal value
    have hGD : ∀ x : Fin N, x.val ≠ 0 →
        g.toFun p₀ (b x) (b x) = D := fun x hx =>
      g_diag_const (by omega) g h_perm x hx
    -- N = 2 (no off-diagonal pairs with both nonzero) vs N ≥ 3
    rcases lt_or_ge N 3 with hN2 | hN3
    · -- N = 2: only nonzero pair is x = j = ⟨1,⋯⟩
      have hNeq : N = 2 := Nat.le_antisymm (Nat.lt_succ_iff.mp hN2) hN
      subst hNeq
      simp only [Fin.sum_univ_two]
      -- simp_rw unfolded p₀ → uniformDist 2 ⋯ in goal; annotate type explicitly
      simp only [
        show g.toFun (uniformDist 2 (by linarith)) (b (0 : Fin 2)) (b (0 : Fin 2)) = 0
          from hG0 0 0 (Or.inl rfl),
        show g.toFun (uniformDist 2 (by linarith)) (b (0 : Fin 2)) (b (1 : Fin 2)) = 0
          from hG0 0 1 (Or.inl rfl),
        show g.toFun (uniformDist 2 (by linarith)) (b (1 : Fin 2)) (b (0 : Fin 2)) = 0
          from hG0 1 0 (Or.inr rfl),
        show g.toFun (uniformDist 2 (by linarith)) (b (1 : Fin 2)) (b (1 : Fin 2)) = D
          from hGD 1 (by decide),
        mul_zero, zero_mul, add_zero, zero_add]
      -- Goal: u 1 * (v 1 * D) = D / 2 * (u 0 * v 0 + u 1 * v 1)
      have hu0 : u 0 = -u 1 := by
        have := hu; simp only [Fin.sum_univ_two] at this; linarith
      have hv0 : v 0 = -v 1 := by
        have := hv; simp only [Fin.sum_univ_two] at this; linarith
      rw [hu0, hv0]; ring
    · -- N ≥ 3: off-diagonal pairs both contribute D/2
      have hGOff : ∀ x j : Fin N, x.val ≠ 0 → j.val ≠ 0 → x ≠ j →
          g.toFun p₀ (b x) (b j) = D / 2 := fun x j hx hj hxj =>
        g_offdiag_half hN3 g h_perm x j hx hj (Fin.val_ne_iff.mpr hxj)
      -- Expand g(u,v) = Σ_{x,j} u_x v_j g(b_x, b_j) using bilinearity
      rw [g_sum_left, g_sum_right]
      -- Partition the sum using b_sum_left/right and g properties
      -- Key: Σ_{x,j: x≠0, j≠0} u_x v_j = u_0 v_0 when Σ u = Σ v = 0
      -- Proof: u_0 = -Σ_{x≠0} u_x (from Σ u = 0), similarly v_0 = -Σ_{j≠0} v_j
      -- So u_0 v_0 = (Σ_{x≠0} u_x)(Σ_{j≠0} v_j) = Σ_{x≠0, j≠0} u_x v_j
      have h_partition : (∑ x : Fin N, ∑ j : Fin N, u x * v j * g.toFun p₀ (b x) (b j)) =
          D * (∑ i : Fin N, u i * v i) / 2 := by
        -- Use sum_erase to remove x=0 and j=0 terms, then apply hGOff for off-diagonal
        sorry
      rw [h_partition]
      
end UniformMetric

-- ============================================================
-- §7  EQUAL REFINEMENTS: CONSTANT IS DIMENSION-INDEPENDENT
-- ============================================================

section RefinementConstant

/-- Equal refinement: split each state into m equal substates.
    This sends uniform_N to uniform_{Nm}. -/
lemma equal_refinement_const {N m : ℕ} (hN : N ≥ 2) (hm : m ≥ 1)
    (g : ∀ n, RiemannianMetric n)
    (h_inv : ∀ n, IsChentsovInvariant (g n) (g (n+1)))
    (h_perm : ∀ n, IsPermutationInvariant (g n)) :
    ∃ C : ℝ, C > 0 ∧
    ∀ N hN, (metric_at_uniform hN (g N) (h_perm N)).choose = C * N := by
  -- Key relation: λ_{Nm} = m · λ_N
  -- Proof: embed uniform_N → uniform_{Nm} by splitting each state into m equal parts.
  -- By invariance: λ_N ∑ u_i v_i = λ_{Nm} · (1/m) ∑ u_i v_i
  -- Hence λ_{Nm} = m · λ_N.
  -- Therefore C = λ_N / N is independent of N.
  sorry -- TODO: formalize the refinement argument

end RefinementConstant

-- ============================================================
-- §8  RATIONAL POINTS: g_p = C · fisherMetric
-- ============================================================

section RationalPoints

/-- For rational p = (k_1/K, ..., k_N/K), refine state i into k_i
    equal substates. This sends p to the uniform distribution on K points.
    By invariance: g_p(u,v) = g_uniform_K(dΦ u, dΦ v).
    Since dΦ u is block-constant with blocks of size k_i,
    g_p(u,v) = C · ∑_i k_i · (u_i/k_i)(v_i/k_i) / (1/K)
             = C · ∑_i u_i v_i / p_i. -/
lemma fisher_on_rational {N : ℕ} (hN : N ≥ 2)
    (g : ∀ n, RiemannianMetric n)
    (h_inv : ∀ n, IsChentsovInvariant (g n) (g (n+1)))
    (h_perm : ∀ n, IsPermutationInvariant (g n)) :
    ∃ C : ℝ, C > 0 ∧
    ∀ (p : openSimplex N) (hp_rat : ∀ i, ∃ k : ℕ, p.1 i = k / (∑ j, (fun j => (Nat.ceil (p.1 j * 1000000) : ℝ)) j)),
    ∀ u v : Fin N → ℝ, ∑ i, u i = 0 → ∑ i, v i = 0 →
      (g N).toFun p u v = C * fisherMetric p u v := by
  sorry -- TODO: formalize rational point argument

end RationalPoints

-- ============================================================
-- §9  MAIN THEOREM: CHENTSOV'S THEOREM
-- ============================================================

section ChentsovTheorem

/-- **Chentsov's theorem.** Any Riemannian metric on Δⁿ (n ≥ 3)
    that is invariant under all Markov embeddings and under permutations,
    and is smooth, must be a positive multiple of the Fisher metric.

    Proof structure:
    1. Permutation invariance → metric = λ_N · Euclidean at uniform point
    2. Equal refinements → λ_{Nm} = mλ_N → C = λ_N/N is dimension-independent
    3. Rational points → refine to uniform → g_p = C · fisherMetric
    4. Density + smoothness → extends to all p
    5. Positivity → C > 0 -/
theorem chentsov_theorem (n : ℕ) (hn : n ≥ 3) (g : RiemannianMetric n)
    (h_inv : IsChentsovInvariant g sorry)
    (h_perm : IsPermutationInvariant g)
    (h_smooth : True) :
    ∃ (c : ℝ), c > 0 ∧ ∀ (p : openSimplex n) (X Y : Fin n → ℝ),
      (∑ i, X i = 0) → (∑ i, Y i = 0) →
      g.toFun p X Y = c * fisherMetric p X Y := by
  sorry

end ChentsovTheorem