import Mathlib.Algebra.Order.BigOperators.Ring.Finset
import Mathlib.Data.Real.Basic
import Mathlib.Tactic.FieldSimp
import Mathlib.Tactic.Ring
import Mathlib.Tactic.Linarith

/-! Samuelson (1968), How Deviant Can You Be?
The single-observation bound for a finite deterministic population.
Variance uses divisor N, not the unbiased sample-estimator divisor N - 1. -/
namespace PaperToLean.Samuelson
noncomputable section
open scoped BigOperators
variable {ι : Type*}

def mean (I : Finset ι) (x : ι → ℝ) : ℝ := (∑ i ∈ I, x i) / I.card
def variance (I : Finset ι) (x : ι → ℝ) : ℝ :=
  (∑ i ∈ I, (x i - mean I x) ^ 2) / I.card

theorem centered_sum (I : Finset ι) (x : ι → ℝ) (hI : I.Nonempty) :
    ∑ i ∈ I, (x i - mean I x) = 0 := by
  have hn : (I.card : ℝ) ≠ 0 := by exact_mod_cast (ne_of_gt hI.card_pos)
  simp only [Finset.sum_sub_distrib, Finset.sum_const, nsmul_eq_mul, mean]
  field_simp
  ring

theorem zero_sum_bound (I : Finset ι) (d : ι → ℝ) (k : ι) (hk : k ∈ I)
    (hd : ∑ i ∈ I, d i = 0) :
    (I.card : ℝ) * (d k) ^ 2 ≤ (I.card - 1 : ℝ) * ∑ i ∈ I, (d i) ^ 2 := by
  classical
  have hc := Finset.sum_mul_sq_le_sq_mul_sq (I.erase k) (fun _ => (1 : ℝ)) d
  simp only [one_mul, one_pow, Finset.sum_const, nsmul_eq_mul, mul_one,
    Finset.sum_erase_eq_sub hk, hd, zero_sub, neg_sq] at hc
  nlinarith

theorem deviation_bound (I : Finset ι) (x : ι → ℝ) (k : ι) (hk : k ∈ I) :
    (x k - mean I x) ^ 2 ≤ (I.card - 1 : ℝ) * variance I x := by
  have hn : (0 : ℝ) < I.card := by exact_mod_cast (Finset.card_pos.mpr ⟨k, hk⟩)
  have h := zero_sum_bound I (fun i => x i - mean I x) k hk (centered_sum I x ⟨k, hk⟩)
  unfold variance
  rw [← mul_div_assoc]
  apply (le_div_iff₀ hn).mpr
  nlinarith

/-- If all other centered observations agree, the bound is an equality. -/
theorem equality_when_others_equal (I : Finset ι) (d : ι → ℝ) (k : ι)
    (hk : k ∈ I) (c : ℝ) (hothers : ∀ i ∈ I, i ≠ k → d i = c)
    (hd : ∑ i ∈ I, d i = 0) :
    (I.card : ℝ) * (d k) ^ 2 = (I.card - 1 : ℝ) * ∑ i ∈ I, (d i) ^ 2 := by
  classical
  have hn : 1 ≤ I.card := Nat.succ_le_of_lt (Finset.card_pos.mpr ⟨k, hk⟩)
  have he : ∑ i ∈ I.erase k, d i = (I.card - 1 : ℝ) * c := by
    rw [Finset.sum_congr rfl (fun i hi => hothers i (Finset.mem_erase.mp hi).2 (Finset.mem_erase.mp hi).1)]
    simp [Finset.card_erase_of_mem hk, Nat.cast_sub hn]
  have he2 : ∑ i ∈ I.erase k, (d i) ^ 2 = (I.card - 1 : ℝ) * c ^ 2 := by
    rw [Finset.sum_congr rfl (fun i hi => congrArg (· ^ 2)
      (hothers i (Finset.mem_erase.mp hi).2 (Finset.mem_erase.mp hi).1))]
    simp [Finset.card_erase_of_mem hk, Nat.cast_sub hn]
  have hz := Finset.sum_erase_add I d hk
  rw [he, hd] at hz
  have hs := Finset.sum_erase_add I (fun i => (d i) ^ 2) hk
  rw [he2] at hs
  have hdk : d k = -((I.card - 1 : ℝ) * c) := by linarith
  rw [← hs, hdk]
  ring

#print axioms centered_sum
#print axioms zero_sum_bound
#print axioms deviation_bound
#print axioms equality_when_others_equal
end
end PaperToLean.Samuelson