import Mathlib.Algebra.BigOperators.Ring.List
import Mathlib.Data.Real.Basic
import Mathlib.Tactic.FieldSimp
import Mathlib.Tactic.Ring
import Mathlib.Tactic.Linarith
import Mathlib.Tactic.Positivity

/-! Welford (1962): exact correctness of incremental mean and squared scatter.
The list computation is a fold from the tail; update can process new samples
online. Exact real arithmetic only, with no roundoff guarantee. -/
namespace PaperToLean.Welford
noncomputable section
structure State where
  count : ℕ
  mean : ℝ
  m2 : ℝ

def update (s : State) (x : ℝ) : State :=
  let μ := s.mean + (x - s.mean) / (s.count + 1 : ℝ)
  ⟨s.count + 1, μ, s.m2 + (x - s.mean) * (x - μ)⟩
def compute : List ℝ → State
  | [] => ⟨0, 0, 0⟩
  | x :: xs => update (compute xs) x

theorem mean_accumulator (s : State) (x : ℝ) :
    ((update s x).count : ℝ) * (update s x).mean = (s.count : ℝ) * s.mean + x := by
  have hn : (s.count : ℝ) + 1 ≠ 0 := by positivity
  simp only [update, Nat.cast_add, Nat.cast_one]
  field_simp
  ring

theorem square_accumulator (s : State) (x : ℝ) :
    (update s x).m2 + ((update s x).count : ℝ) * (update s x).mean ^ 2 =
      s.m2 + (s.count : ℝ) * s.mean ^ 2 + x ^ 2 := by
  have hn : (s.count : ℝ) + 1 ≠ 0 := by positivity
  simp only [update, Nat.cast_add, Nat.cast_one]
  field_simp
  ring

theorem scatter_increment (s : State) (x : ℝ) :
    (update s x).m2 = s.m2 + (s.count : ℝ) / (s.count + 1 : ℝ) * (x - s.mean) ^ 2 := by
  have hn : (s.count : ℝ) + 1 ≠ 0 := by positivity
  simp only [update]
  field_simp
  ring

theorem update_nonneg (s : State) (x : ℝ) (hs : 0 ≤ s.m2) : 0 ≤ (update s x).m2 := by
  rw [scatter_increment]
  positivity

/-- Both moment invariants and the count are derived from the computation. -/
theorem moments (xs : List ℝ) :
    (compute xs).count = xs.length ∧
    ((compute xs).count : ℝ) * (compute xs).mean = xs.sum ∧
    (compute xs).m2 + ((compute xs).count : ℝ) * (compute xs).mean ^ 2 =
      (xs.map (· ^ 2)).sum := by
  induction xs with
  | nil => simp [compute]
  | cons x xs ih =>
    refine ⟨by simp [compute, update, ih.1], ?_, ?_⟩
    · rw [compute, mean_accumulator, ih.2.1, List.sum_cons]
      ring
    · rw [compute, square_accumulator, ih.2.2, List.map_cons, List.sum_cons]
      ring

theorem mean_correct (xs : List ℝ) (hn : xs ≠ []) :
    (compute xs).mean = xs.sum / xs.length := by
  have hp : 0 < xs.length := List.length_pos_iff.mpr hn
  have hne : (xs.length : ℝ) ≠ 0 := by exact_mod_cast (ne_of_gt hp)
  apply (eq_div_iff hne).mpr
  have h := (moments xs).2.1
  rw [(moments xs).1] at h
  nlinarith

theorem centered_sum (xs : List ℝ) (μ : ℝ) :
    (xs.map (fun x => (x - μ) ^ 2)).sum =
      (xs.map (· ^ 2)).sum - 2 * μ * xs.sum + (xs.length : ℝ) * μ ^ 2 := by
  induction xs with
  | nil => simp
  | cons x xs ih =>
    simp only [List.map_cons, List.sum_cons, List.length_cons, Nat.cast_add, Nat.cast_one, ih]
    ring

theorem scatter_correct (xs : List ℝ) :
    (compute xs).m2 = (xs.map (fun x => (x - (compute xs).mean) ^ 2)).sum := by
  obtain ⟨hc, hμ, hM⟩ := moments xs
  rw [centered_sum, ← hM, ← hμ, ← hc]
  ring

#print axioms mean_accumulator
#print axioms square_accumulator
#print axioms scatter_increment
#print axioms update_nonneg
#print axioms moments
#print axioms mean_correct
#print axioms centered_sum
#print axioms scatter_correct
end
end PaperToLean.Welford