-- CvxLean/Lib/Math/CovarianceEstimation.lean

import CvxLean.Lib.Math.LinearAlgebra.Matrix.PosDef

noncomputable section CovarianceEstimation

open Real

open Matrix

open BigOperators

def Matrix.quadForm {n : ℕ} (R : Matrix (Fin n) (Fin n) ℝ) (x : Fin n → ℝ) : ℝ :=
  x ⬝ᵥ R.mulVec x

def gaussianPdf {n : ℕ} (R : Matrix (Fin n) (Fin n) ℝ) (x : Fin n → ℝ) : ℝ :=
  (1 / sqrt (((2 * π) ^ n) * R.det)) * exp (- R⁻¹.quadForm x / 2)

-- Repository definitions used:
content: def mulVec (M : Matrix (Fin n) (Fin m) Float) (v : (Fin m) → Float) : Fin n → Float :=
  fun i => (fun j => M i j) ⬝ᵥᶜ v

content: def dotProduct (v w : Fin n → Float) : Float :=
  (Array.zipWith (Vec.Computable.toArray v) (Vec.Computable.toArray w) Float.mul).foldl Float.add 0

content: def toArray (x : Fin n → Float) : Array Float :=
  (Array.range n).map (fun i => if h : i < n then x ⟨i, h⟩ else 0)

content: def det {n : ℕ} (A : Matrix (Fin n) (Fin n) Float) : Float :=
  if h : 0 < n then
    if n == 1 then A ⟨0, h⟩ ⟨0, h⟩ else
      (List.finRange n).foldl (fun s i =>
        s + (-1) ^ (Float.ofNat i.val) * A i ⟨0, h⟩ * det (minor A i ⟨0, h⟩)) 0
  else 0

content: def minor (A : Matrix (Fin n) (Fin n) Float) (a b : Fin n) :
    Matrix (Fin n.pred) (Fin n.pred) Float :=
  match n with
  | 0 => fun _ => Fin.elim0
  | _ + 1 => minorAux A a b

content: private def minorAux (A : Matrix (Fin n.succ) (Fin n.succ) Float) (a b : Fin n.succ) :
    Matrix (Fin n) (Fin n) Float :=
  fun i j =>
    let i' : Fin n.succ := if i.val < a.val then i else i.succ;
    let j' : Fin n.succ := if j.val < b.val then j else j.succ;
    A i' j'

content: noncomputable instance PosDef.Invertible [DecidableEq n] {M : Matrix n n 𝕜} (hM : M.PosDef) :
    Invertible M :=
  invertibleOfIsUnitDet M (isUnit_iff_ne_zero.2 hM.det_ne_zero)

content: infixl:72 " ⬝ᵥᶜ " => Matrix.Computable.dotProduct

-- Repository lemmas used:
content: theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) :\n(l₁ ++ l₂)[i] = l₁[i]

-- Local definitions used:
content: def Matrix.quadForm {n : ℕ} (R : Matrix (Fin n) (Fin n) ℝ) (x : Fin n → ℝ) : ℝ :=
  x ⬝ᵥ R.mulVec x

content: def gaussianPdf {n : ℕ} (R : Matrix (Fin n) (Fin n) ℝ) (x : Fin n → ℝ) : ℝ :=
  (1 / sqrt (((2 * π) ^ n) * R.det)) * exp (- R⁻¹.quadForm x / 2)

-- Local lemmas used:
content: lemma gaussianPdf_pos {n : ℕ} (R : Matrix (Fin n) (Fin n) ℝ) (y : Fin n → ℝ) (h : R.PosDef) :
    0 < gaussianPdf R y

-- Target theorem and proof:
lemma log_prod_gaussianPdf {N n : ℕ} (y : Fin N → Fin n → ℝ) (R : Matrix (Fin n) (Fin n) ℝ)
    (hR : R.PosDef) : log (∏ i : Fin N, gaussianPdf R (y i)) =
    ∑ i : Fin N, (- (log (sqrt ((2 * π) ^ n)) + log (sqrt R.det)) + - R⁻¹.quadForm (y i) / 2)
:= by
  have : ∀ i,
    i ∈ Finset.univ → gaussianPdf R (y i) ≠ 0 := fun i _ => ne_of_gt (gaussianPdf_pos _ _ hR)
  have sqrt_2_pi_n_R_det_ne_zero: sqrt ((2 * π) ^ n * R.det) ≠ 0 := by
    refine' ne_of_gt (sqrt_pos.2 (mul_pos _ hR.det_pos))
    exact (pow_pos (mul_pos (by positivity) pi_pos) _)
  rw [log_prod Finset.univ (fun i => gaussianPdf R (y i)) this]
  unfold gaussianPdf
  apply congr_arg (Finset.sum Finset.univ)
  ext i
  rw [log_mul, log_div, sqrt_mul, log_mul, log_exp, log_one, zero_sub]
  simp [rpow_eq_pow]
  exact ne_of_gt (sqrt_pos.2 (pow_pos (mul_pos (by positivity) pi_pos) _))
  exact ne_of_gt (sqrt_pos.2 hR.det_pos)
  exact pow_nonneg (mul_nonneg (by positivity) (le_of_lt pi_pos)) _
  norm_num
  exact sqrt_2_pi_n_R_det_ne_zero
  exact div_ne_zero (by norm_num) sqrt_2_pi_n_R_det_ne_zero
  exact exp_ne_zero _