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

/-! Grüss (1935): a finite weighted analogue of the integral covariance bound.
Nonnegative weights sum to one. This file does not formalize integration.
Finite Cauchy--Schwarz is reused from Mathlib, not claimed as a new result. -/
namespace PaperToLean.Gruss
noncomputable section
open scoped BigOperators
variable {ι : Type*}
def avg (I : Finset ι) (w f : ι → ℝ) : ℝ := ∑ i ∈ I, w i * f i
def variance (I : Finset ι) (w f : ι → ℝ) : ℝ := avg I w (fun i => f i ^ 2) - (avg I w f) ^ 2
def covariance (I : Finset ι) (w f g : ι → ℝ) : ℝ := avg I w (fun i => f i * g i) - avg I w f * avg I w g

theorem centered_covariance (I : Finset ι) (w f g : ι → ℝ) (hw : ∑ i ∈ I, w i = 1) :
    avg I w (fun i => (f i - avg I w f) * (g i - avg I w g)) = covariance I w f g := by
  have hexp (u v : ℝ) : (∑ i ∈ I, w i * ((f i - u) * (g i - v))) =
      avg I w (fun i => f i * g i) - v * avg I w f -
      u * avg I w g + u * v * ∑ i ∈ I, w i := by
    simp only [avg, Finset.mul_sum, ← Finset.sum_sub_distrib,
      ← Finset.sum_add_distrib]
    apply Finset.sum_congr rfl
    intro i _
    ring
  change (∑ i ∈ I, w i * ((f i - avg I w f) * (g i - avg I w g))) = _
  rw [hexp, hw]
  unfold covariance
  ring

theorem centered_variance (I : Finset ι) (w f : ι → ℝ) (hw : ∑ i ∈ I, w i = 1) :
    avg I w (fun i => (f i - avg I w f) ^ 2) = variance I w f := by
  simpa only [variance, covariance, pow_two] using centered_covariance I w f f hw

theorem variance_nonneg (I : Finset ι) (w f : ι → ℝ)
    (hw : ∑ i ∈ I, w i = 1) (hw0 : ∀ i ∈ I, 0 ≤ w i) : 0 ≤ variance I w f := by
  rw [← centered_variance I w f hw]
  exact Finset.sum_nonneg (fun i hi => mul_nonneg (hw0 i hi) (sq_nonneg _))

theorem range_variance_bound (I : Finset ι) (w f : ι → ℝ) (m M : ℝ)
    (hw : ∑ i ∈ I, w i = 1) (hw0 : ∀ i ∈ I, 0 ≤ w i)
    (hlo : ∀ i ∈ I, m ≤ f i) (hhi : ∀ i ∈ I, f i ≤ M) :
    variance I w f ≤ (M - m) ^ 2 / 4 := by
  have hpoint : ∀ i ∈ I, w i * f i ^ 2 ≤ w i * ((m + M) * f i - m * M) := by
    intro i hi
    apply mul_le_mul_of_nonneg_left _ (hw0 i hi)
    nlinarith [mul_nonneg (sub_nonneg.mpr (hlo i hi)) (sub_nonneg.mpr (hhi i hi))]
  have hs := Finset.sum_le_sum hpoint
  have heq : (∑ i ∈ I, w i * ((m + M) * f i - m * M)) =
      (m + M) * avg I w f - m * M := by
    calc
      _ = (m + M) * avg I w f - m * M * (∑ i ∈ I, w i) := by
        simp only [avg, Finset.mul_sum, ← Finset.sum_sub_distrib]
        apply Finset.sum_congr rfl
        intro i _
        ring
      _ = _ := by rw [hw, mul_one]
  rw [heq] at hs
  change avg I w (fun i => f i ^ 2) ≤ _ at hs
  unfold variance
  nlinarith [sq_nonneg (avg I w f - (m + M) / 2)]

theorem covariance_cauchy (I : Finset ι) (w f g : ι → ℝ)
    (hw : ∑ i ∈ I, w i = 1) (hw0 : ∀ i ∈ I, 0 ≤ w i) :
    (covariance I w f g) ^ 2 ≤ variance I w f * variance I w g := by
  rw [← centered_covariance I w f g hw, ← centered_variance I w f hw,
    ← centered_variance I w g hw]
  apply Finset.sum_sq_le_sum_mul_sum_of_sq_le_mul I
  · intro i hi
    exact mul_nonneg (hw0 i hi) (sq_nonneg _)
  · intro i hi
    exact mul_nonneg (hw0 i hi) (sq_nonneg _)
  · intro i _
    exact le_of_eq (by ring)

theorem inequality (I : Finset ι) (w f g : ι → ℝ) (m M p P : ℝ)
    (hw : ∑ i ∈ I, w i = 1) (hw0 : ∀ i ∈ I, 0 ≤ w i)
    (hmM : m ≤ M) (hpP : p ≤ P)
    (hf : ∀ i ∈ I, m ≤ f i ∧ f i ≤ M) (hg : ∀ i ∈ I, p ≤ g i ∧ g i ≤ P) :
    |covariance I w f g| ≤ (M - m) * (P - p) / 4 := by
  have hc := covariance_cauchy I w f g hw hw0
  have hfv := range_variance_bound I w f m M hw hw0 (fun i hi => (hf i hi).1) (fun i hi => (hf i hi).2)
  have hgv := range_variance_bound I w g p P hw hw0 (fun i hi => (hg i hi).1) (fun i hi => (hg i hi).2)
  have hprod := mul_le_mul hfv hgv (variance_nonneg I w g hw hw0) (by positivity : 0 ≤ (M - m) ^ 2 / 4)
  have hb : 0 ≤ (M - m) * (P - p) / 4 := by positivity
  have heq : ((M - m) * (P - p) / 4) ^ 2 = ((M - m) ^ 2 / 4) * ((P - p) ^ 2 / 4) := by ring
  have habs := sq_abs (covariance I w f g)
  nlinarith [abs_nonneg (covariance I w f g)]

#print axioms centered_covariance
#print axioms centered_variance
#print axioms variance_nonneg
#print axioms range_variance_bound
#print axioms covariance_cauchy
#print axioms inequality
end
end PaperToLean.Gruss