module
public import Mathlib.LinearAlgebra.Matrix.Trace
public import Mathlib.LinearAlgebra.Matrix.Symmetric
public import Mathlib.Tactic

/-!
Real, finite-basis closed-shell density identities. `C` contains occupied
spatial orbitals; the density includes both spins. No SCF convergence,
floating-point arithmetic, or continuum integral evaluator is assumed proved.
-/

@[expose] public section

noncomputable section

namespace VibeQCRHF

open Matrix

variable {AO Occ : Type*}

def density [Fintype Occ] (C : Matrix AO Occ ℝ) : Matrix AO AO ℝ :=
  (2 : ℝ) • (C * Cᵀ)

theorem density_symmetric [Fintype Occ] (C : Matrix AO Occ ℝ) :
    (density C)ᵀ = density C := by
  simp [density, transpose_mul]

theorem density_metric_idempotent [Fintype AO] [Fintype Occ] [DecidableEq Occ] (C : Matrix AO Occ ℝ) (S : Matrix AO AO ℝ)
    (hC : Cᵀ * S * C = 1) :
    density C * S * density C = (2 : ℝ) • density C := by
  have h : (C * Cᵀ) * S * (C * Cᵀ) = C * Cᵀ := by
    calc
      _ = C * (Cᵀ * S * C) * Cᵀ := by simp only [Matrix.mul_assoc]
      _ = _ := by rw [hC]; simp
  simp only [density, Matrix.smul_mul, Matrix.mul_smul, smul_smul]
  rw [h]

theorem density_electron_count [Fintype AO] [Fintype Occ] [DecidableEq Occ] (C : Matrix AO Occ ℝ) (S : Matrix AO AO ℝ)
    (hC : Cᵀ * S * C = 1) :
    Matrix.trace (density C * S) = 2 * (Fintype.card Occ : ℝ) := by
  have h : Matrix.trace ((C * Cᵀ) * S) = (Fintype.card Occ : ℝ) := by
    calc
      _ = Matrix.trace (Cᵀ * S * C) := by
        rw [Matrix.mul_assoc, Matrix.trace_mul_comm]
      _ = _ := by rw [hC]; simp
  simpa [density, Matrix.smul_mul, Matrix.trace_smul] using congrArg (fun x : ℝ => 2 * x) h

theorem density_occupied_rotation [Fintype Occ] [DecidableEq Occ] (C : Matrix AO Occ ℝ) (U : Matrix Occ Occ ℝ)
    (hU : U * Uᵀ = 1) : density (C * U) = density C := by
  unfold density
  congr 1
  rw [transpose_mul]
  calc
    _ = C * (U * Uᵀ) * Cᵀ := by simp only [Matrix.mul_assoc]
    _ = _ := by rw [hU]; simp

/-- A generalized occupied eigenpair implies a zero metric commutator.
The symmetric multiplier can be diagonal orbital energies. This is a
stationarity identity, not a convergence or minimum-energy theorem. -/
theorem roothaan_implies_stationarity [Fintype AO] [Fintype Occ]
    (C : Matrix AO Occ ℝ) (S F : Matrix AO AO ℝ) (eps : Matrix Occ Occ ℝ)
    (hS : Sᵀ = S) (hF : Fᵀ = F) (heps : epsᵀ = eps)
    (heig : F * C = S * C * eps) :
    F * density C * S = S * density C * F := by
  have ht : Cᵀ * F = eps * Cᵀ * S := by
    have ht := congrArg Matrix.transpose heig
    simpa [transpose_mul, hS, hF, heps, Matrix.mul_assoc] using ht
  have h : F * (C * Cᵀ) * S = S * (C * Cᵀ) * F := by
    calc
      _ = (F * C) * Cᵀ * S := by simp only [Matrix.mul_assoc]
      _ = (S * C * eps) * Cᵀ * S := by rw [heig]
      _ = S * C * (eps * Cᵀ * S) := by simp only [Matrix.mul_assoc]
      _ = S * C * (Cᵀ * F) := by rw [ht]
      _ = _ := by simp only [Matrix.mul_assoc]
  simpa only [density, Matrix.smul_mul, Matrix.mul_smul] using congrArg (fun A => (2 : ℝ) • A) h

end VibeQCRHF