/-
# Flow-matching map vs. Bures–Wasserstein OT map — the commuting case (Lean anchor)

This file is the **fidelity anchor** for the informal theorem in
`paper/theorem.md` (§5, §9). It formalises the *commuting* case via the
**basis-free continuous-functional-calculus route** (the route the frame
mandates: robust under repeated eigenvalues, no eigenbasis choice).

`Matrix n n ℝ` is a C⋆-algebra, so Mathlib's `CFC.sqrt` and the Loewner order
(`scoped MatrixOrder`) are available. For SPD `σ₀, σ₁` we work entirely inside
the commutative ⋆-subalgebra generated by `σ₀, σ₁` — exactly the algebra `𝒜` of
the informal proof.

## Statement ↔ paper map (constitution pillar: "Fidelity audit")

| Lean                                                | paper/theorem.md                                  |
|-----------------------------------------------------|---------------------------------------------------|
| `flowMap σ₀ σ₁ = √σ₁ * (√σ₀)⁻¹`                     | §5: `T₁ = σ₁^{1/2} σ₀^{-1/2}` (commuting)         |
| `otMap σ₀ σ₁ = (√σ₀)⁻¹ √(√σ₀ σ₁ √σ₀) (√σ₀)⁻¹`      | §0/§1: `T_OT = σ₀^{-1/2}(σ₀^{1/2}σ₁σ₀^{1/2})^{1/2}σ₀^{-1/2}` |
| `(flowMap …).IsHermitian`                           | clause (2): `T₁` symmetric                        |
| `flowMap * σ₀ * (flowMap)ᴴ = σ₁`                    | Lemma L_B (keystone): Bures pushforward           |
| `flowMap = otMap`                                   | clause (3): `T₁ = T_OT`                           |

The matrix-ODE identification `Φ₁ = flowMap` (that the flow map of `(⋆)` equals
this algebraic `flowMap`) and the non-commuting converse are the disclosed
infrastructure gaps; see `lean/unproved.md`. They are **not** axiomatised here.
-/
import Mathlib

open scoped MatrixOrder Matrix

namespace FMG

variable {n : Type*} [Fintype n] [DecidableEq n]

/-- The flow-matching transport map in the commuting case: `√σ₁ · (√σ₀)⁻¹`
(`paper/theorem.md` §5, `T₁ = σ₁^{1/2} σ₀^{-1/2}`). -/
noncomputable def flowMap (σ₀ σ₁ : Matrix n n ℝ) : Matrix n n ℝ :=
  CFC.sqrt σ₁ * (CFC.sqrt σ₀)⁻¹

/-- The Bures–Wasserstein optimal-transport map
`σ₀^{-1/2}(σ₀^{1/2}σ₁σ₀^{1/2})^{1/2}σ₀^{-1/2}` (`paper/theorem.md` §0). -/
noncomputable def otMap (σ₀ σ₁ : Matrix n n ℝ) : Matrix n n ℝ :=
  (CFC.sqrt σ₀)⁻¹ * CFC.sqrt (CFC.sqrt σ₀ * σ₁ * CFC.sqrt σ₀) * (CFC.sqrt σ₀)⁻¹

/-! ### Infrastructure lemmas -/

set_option linter.unusedDecidableInType false in
/-- The CFC square root commutes with anything that commutes with its argument
(the algebraic content of "everything lives in one commutative ⋆-subalgebra"). -/
theorem commute_sqrt_left {a b : Matrix n n ℝ} (h : Commute a b) :
    Commute (CFC.sqrt a) b := by
  rw [CFC.sqrt_eq_cfc]
  exact h.cfc_nnreal NNReal.sqrt

/-- For an invertible matrix, `A⁻¹` commutes with whatever `A` commutes with. -/
theorem commute_inv_left {A B : Matrix n n ℝ} (hA : IsUnit A) (h : Commute A B) :
    Commute A⁻¹ B := by
  have hd : IsUnit A.det := (Matrix.isUnit_iff_isUnit_det A).mp hA
  have hl : A⁻¹ * A = 1 := Matrix.nonsing_inv_mul A hd
  have hr : A * A⁻¹ = 1 := Matrix.mul_nonsing_inv A hd
  have h' : A * B = B * A := h.eq
  calc A⁻¹ * B
      = A⁻¹ * B * (A * A⁻¹) := by rw [hr, mul_one]
    _ = A⁻¹ * (B * A) * A⁻¹ := by simp only [mul_assoc]
    _ = A⁻¹ * (A * B) * A⁻¹ := by rw [← h']
    _ = A⁻¹ * A * B * A⁻¹ := by simp only [mul_assoc]
    _ = B * A⁻¹ := by rw [hl, one_mul]

/-! ### The commuting-case theorem (fidelity anchor) -/

/-- **Flow matching = optimal transport in the commuting case.**

For SPD `σ₀, σ₁` with `[σ₀,σ₁]=0`, the flow-matching map `flowMap σ₀ σ₁ =
√σ₁(√σ₀)⁻¹` is

1. Hermitian (symmetric);
2. a genuine pushforward `N(0,σ₀) → N(0,σ₁)` — it solves the Bures equation
   `T σ₀ Tᴴ = σ₁` (Lemma L_B keystone);
3. equal to the Bures–Wasserstein OT map `otMap σ₀ σ₁`.

This is the Lean image of `paper/theorem.md` clauses (1)⇒(4)⇒(3)⇒(2) in the
commuting direction. -/
theorem flowMap_isHermitian_pushforward_eq_otMap_of_commute
    {σ₀ σ₁ : Matrix n n ℝ} (h0 : σ₀.PosDef) (h1 : σ₁.PosDef)
    (hcomm : Commute σ₀ σ₁) :
    (flowMap σ₀ σ₁).IsHermitian ∧
      flowMap σ₀ σ₁ * σ₀ * (flowMap σ₀ σ₁)ᴴ = σ₁ ∧
      flowMap σ₀ σ₁ = otMap σ₀ σ₁ := by
  simp only [flowMap, otMap]
  set S0 := CFC.sqrt σ₀ with hS0def
  set S1 := CFC.sqrt σ₁ with hS1def
  -- positivity facts
  have h0' : (0 : Matrix n n ℝ) ≤ σ₀ := h0.posSemidef.nonneg
  have h1' : (0 : Matrix n n ℝ) ≤ σ₁ := h1.posSemidef.nonneg
  have hS0nn : (0 : Matrix n n ℝ) ≤ S0 := CFC.sqrt_nonneg σ₀
  have hS1nn : (0 : Matrix n n ℝ) ≤ S1 := CFC.sqrt_nonneg σ₁
  have hS0pd : S0.PosDef :=
    Matrix.isStrictlyPositive_iff_posDef.mp (IsStrictlyPositive.sqrt σ₀ h0.isStrictlyPositive)
  have hS0unit : IsUnit S0 := hS0pd.isUnit
  have hd : IsUnit S0.det := (Matrix.isUnit_iff_isUnit_det S0).mp hS0unit
  -- square-root and inverse identities
  have hSS0 : S0 * S0 = σ₀ := CFC.sqrt_mul_sqrt_self σ₀ h0'
  have hSS1 : S1 * S1 = σ₁ := CFC.sqrt_mul_sqrt_self σ₁ h1'
  have hl : S0⁻¹ * S0 = 1 := Matrix.nonsing_inv_mul S0 hd
  have hr : S0 * S0⁻¹ = 1 := Matrix.mul_nonsing_inv S0 hd
  -- commutations: everything lives in one commutative ⋆-subalgebra
  have cS0sig1 : Commute S0 σ₁ := commute_sqrt_left hcomm
  have cS1sig0 : Commute S1 σ₀ := commute_sqrt_left hcomm.symm
  have cS0S1 : Commute S0 S1 := commute_sqrt_left cS1sig0.symm
  have cS0iS1 : Commute S0⁻¹ S1 := commute_inv_left hS0unit cS0S1
  -- Hermitian facts
  have hS0herm : S0ᴴ = S0 :=
    (Matrix.isHermitian_iff_isSelfAdjoint.mpr (IsSelfAdjoint.of_nonneg hS0nn)).eq
  have hS1herm : S1ᴴ = S1 :=
    (Matrix.isHermitian_iff_isSelfAdjoint.mpr (IsSelfAdjoint.of_nonneg hS1nn)).eq
  have hS0iherm : (S0⁻¹)ᴴ = S0⁻¹ := by
    rw [Matrix.conjTranspose_nonsing_inv, hS0herm]
  -- (1) Hermitian
  have herm : (S1 * S0⁻¹)ᴴ = S1 * S0⁻¹ := by
    rw [Matrix.conjTranspose_mul, hS0iherm, hS1herm]
    exact cS0iS1.eq
  refine ⟨herm, ?_, ?_⟩
  · -- (2) Bures pushforward  T σ₀ Tᴴ = σ₁
    rw [herm]
    calc S1 * S0⁻¹ * σ₀ * (S1 * S0⁻¹)
        = S1 * S0⁻¹ * (S0 * S0) * (S1 * S0⁻¹) := by rw [hSS0]
      _ = S1 * (S0⁻¹ * S0) * S0 * S1 * S0⁻¹ := by simp only [mul_assoc]
      _ = S1 * S0 * S1 * S0⁻¹ := by rw [hl, mul_one]
      _ = S1 * (S1 * S0) * S0⁻¹ := by rw [mul_assoc S1 S0 S1, cS0S1.eq]
      _ = S1 * S1 * (S0 * S0⁻¹) := by simp only [mul_assoc]
      _ = σ₁ := by rw [hSS1, hr, mul_one]
  · -- (3) flowMap = otMap
    have inner_sqrt : CFC.sqrt (S0 * σ₁ * S0) = S1 * S0 := by
      apply CFC.sqrt_unique
      · calc S1 * S0 * (S1 * S0)
            = S1 * (S0 * S1) * S0 := by simp only [mul_assoc]
          _ = S1 * (S1 * S0) * S0 := by rw [cS0S1.eq]
          _ = S1 * S1 * (S0 * S0) := by simp only [mul_assoc]
          _ = σ₁ * (S0 * S0) := by rw [hSS1]
          _ = S0 * σ₁ * S0 := by
              rw [show S0 * σ₁ = σ₁ * S0 from cS0sig1.eq, mul_assoc]
      · exact Commute.mul_nonneg hS1nn hS0nn cS0S1.symm
    rw [inner_sqrt]
    symm
    calc S0⁻¹ * (S1 * S0) * S0⁻¹
        = S0⁻¹ * S1 * (S0 * S0⁻¹) := by simp only [mul_assoc]
      _ = S0⁻¹ * S1 := by rw [hr, mul_one]
      _ = S1 * S0⁻¹ := cS0iS1.eq

end FMG