-- Variance and Chebyshev's inequality
-- Part of Cumulus, the Grace17 mathematics library.
-- Copyright (c) 2026 Grant Passmore.
-- Released under the Apache License 2.0. See LICENSE in the repository root.
-- Authors: Grant Passmore

module Grace.Cumulus.Probability.Variance
import RandomVariable

-- The pointwise square of a random random.
defgraph squareVariable (Ω X : Set) (x ∈ Ω) : ℝ := call2(realMul, X(x), X(x))
requires random : X : Ω → ℝ
proof
  have value := FunctionBetween.apply_closed with random, member
  exact realMulValue.real at X(x), X(x)

-- A random random shifted by the constant m.
defgraph centeredVariable (Ω X : Set) (m : Real) (x ∈ Ω) : ℝ :=
    realSubValue(X(x), m)
requires random : X : Ω → ℝ
proof
  have value := FunctionBetween.apply_closed with random, member
  exact realSubValue.real at X(x), m

def variance (Ω P X : Set) : Set :=
  expectation(Ω, P, squareVariable(Ω, centeredVariable(Ω, X, expectation(Ω, P,
    X))))

theorem variance.value (Ω P X : Set)
  shows variance(Ω, P, X) =
    expectation(Ω, P, squareVariable(Ω, centeredVariable(Ω, X, expectation(Ω,
      P, X))))
proof
  simp [variance]

theorem squareVariable.nonnegative (Ω X ω : Set)
  requires random : X : Ω → ℝ
  requires point : ω ∈ Ω
  shows (0:R) ≤ (squareVariable(Ω, X)(ω):R)
proof
  have xReal := FunctionBetween.apply_closed with random, point
  have value := squareVariable.value at Ω, X, ω with random, point
  rw value
  exact realMulValue.square_nonnegative at X(ω)

theorem expectation.square_nonnegative (Ω P X : Set)
  requires space : ProbabilitySpace(Ω, P)
  requires random : X : Ω → ℝ
  shows (0:R) ≤ (expectation(Ω, P, squareVariable(Ω, X)):R)
proof
  have square := squareVariable.function with random
  have bounds : ∀ ω ∈ Ω, (0:R) ≤ (squareVariable(Ω, X)(ω):R)
    intros ω point
    exact squareVariable.nonnegative with random, point
  exact expectation.nonnegative with space, square, bounds

theorem variance.real (Ω P X : Set)
  requires space : ProbabilitySpace(Ω, P)
  requires random : X : Ω → ℝ
  shows Real(variance(Ω, P, X))
proof
  have mean := expectation.real with space, random
  have centered := centeredVariable.function at Ω, X, expectation(Ω, P, X) with
      random
  have square := squareVariable.function with centered
  rw variance.value
  exact expectation.real with space, square

theorem variance.nonnegative (Ω P X : Set)
  requires space : ProbabilitySpace(Ω, P)
  requires random : X : Ω → ℝ
  shows (0:R) ≤ (variance(Ω, P, X):R)
proof
  have mean := expectation.real with space, random
  have centered := centeredVariable.function at Ω, X, expectation(Ω, P, X) with
      random
  rw variance.value
  exact expectation.square_nonnegative with space, centered

-- The computational formula: Var X = E[X²] - (E X)².
theorem variance.formula (Ω P X : Set)
  requires space : ProbabilitySpace(Ω, P)
  requires random : X : Ω → ℝ
  shows variance(Ω, P, X) = (expectation(Ω, P, squareVariable(Ω, X)):R) -
    (expectation(Ω, P, X):R) * (expectation(Ω, P, X):R)
proof
  have mean := expectation.real with space, random
  have twoReal := realNumeral.real at 2
  have twiceMean := realMulValue.real at (2:R), expectation(Ω, P, X)
  have squaredMean := realMulValue.real at expectation(Ω, P, X), expectation(Ω,
      P, X)
  have centered := centeredVariable.function at Ω, X, expectation(Ω, P, X) with
      random
  have centeredSquare := squareVariable.function with centered
  have square := squareVariable.function with random
  have scaled := weightScale.function at Ω, (2:R) * (expectation(Ω, P, X):R),
      X with random
  have left := weightAdd.function with centeredSquare, scaled
  have constant := constantFunction.function at Ω, ℝ,
    (expectation(Ω, P, X):R) * (expectation(Ω, P, X):R) with squaredMean
  have right := weightAdd.function with square, constant
  have agree : ∀ ω ∈ Ω, weightAdd(Ω, squareVariable(Ω, centeredVariable(Ω, X,
      expectation(Ω, P, X))),
      weightScale(Ω, (2:R) * (expectation(Ω, P, X):R), X))(ω) =
    weightAdd(Ω, squareVariable(Ω, X), constantFunction(Ω, ℝ,
      (expectation(Ω, P, X):R) * (expectation(Ω, P, X):R)))(ω)
    intros ω point
    have xReal := FunctionBetween.apply_closed with random, point
    have leftValue := weightAdd.value at Ω,
      squareVariable(Ω, centeredVariable(Ω, X, expectation(Ω, P, X))),
      weightScale(Ω, (2:R) * (expectation(Ω, P, X):R), X), ω with
        centeredSquare, scaled, point
    have centeredSquareValue := squareVariable.value at Ω,
      centeredVariable(Ω, X, expectation(Ω, P, X)), ω with centered, point
    have centeredValue := centeredVariable.value at Ω, X, expectation(Ω, P, X),
        ω with
      random, point
    have scaledValue := weightScale.value at Ω, (2:R) * (expectation(Ω, P,
        X):R), X, ω with
      random, point
    have rightValue := weightAdd.value at Ω, squareVariable(Ω, X),
        constantFunction(Ω, ℝ,
      (expectation(Ω, P, X):R) * (expectation(Ω, P, X):R)), ω with square,
        constant, point
    have squareValue := squareVariable.value at Ω, X, ω with random, point
    have constantValue := constantFunction.apply at Ω, ℝ, ω,
      (expectation(Ω, P, X):R) * (expectation(Ω, P, X):R) with point,
        squaredMean
    rw [leftValue, centeredSquareValue, centeredValue, scaledValue, rightValue,
      squareValue,
      constantValue]
    ring
  have same := expectation.congr with space, left, right, agree
  have leftSplit := expectation.add with space, centeredSquare, scaled
  have scale := expectation.scale at Ω, P, X, (2:R) * (expectation(Ω, P,
      X):R) with
    space, random
  have rightSplit := expectation.add with space, square, constant
  have constantValue := expectation.constant at Ω, P,
    (expectation(Ω, P, X):R) * (expectation(Ω, P, X):R) with space
  rw [leftSplit, scale, rightSplit, constantValue] at same
  have varianceReal := variance.real with space, random
  have squareReal := expectation.real with space, square
  rw variance.value at varianceReal
  rw variance.value
  calc
    expectation(Ω, P, squareVariable(Ω, centeredVariable(Ω, X, expectation(Ω,
      P, X)))) =
        ((expectation(Ω, P, squareVariable(Ω, centeredVariable(Ω, X,
          expectation(Ω, P, X)))):R) +
          ((2:R) * (expectation(Ω, P, X):R)) * (expectation(Ω, P, X):R)) -
          ((2:R) * (expectation(Ω, P, X):R)) * (expectation(Ω, P,
            X):R) := by ring
    _ = ((expectation(Ω, P, squareVariable(Ω, X)):R) +
          (expectation(Ω, P, X):R) * (expectation(Ω, P, X):R)) -
          ((2:R) * (expectation(Ω, P, X):R)) * (expectation(Ω, P, X):R) := by cc
    _ = (expectation(Ω, P, squareVariable(Ω, X)):R) -
          (expectation(Ω, P, X):R) * (expectation(Ω, P, X):R) := by ring

theorem variance.le_square_expectation (Ω P X : Set)
  requires space : ProbabilitySpace(Ω, P)
  requires random : X : Ω → ℝ
  shows (variance(Ω, P, X):R) ≤ (expectation(Ω, P, squareVariable(Ω, X)):R)
proof
  have mean := expectation.real with space, random
  have square := squareVariable.function with random
  have squareReal := expectation.real with space, square
  have varianceReal := variance.real with space, random
  have formula := variance.formula with space, random
  have nonnegative := realMulValue.square_nonnegative at expectation(Ω, P, X)
  rw formula
  linarith

-- Scaling a random by c scales its variance by c².
theorem variance.scale (Ω P X : Set) (c : Real)
  requires space : ProbabilitySpace(Ω, P)
  requires random : X : Ω → ℝ
  shows variance(Ω, P, weightScale(Ω, c, X)) = (c * c) * (variance(Ω, P, X):R)
proof
  have mean := expectation.real with space, random
  have cSquare := realMulValue.real at c, c
  have scaled := weightScale.function at Ω, c, X with random
  have scaledMean := expectation.scale at Ω, P, X, c with space, random
  have centered := centeredVariable.function at Ω, X, expectation(Ω, P, X) with
      random
  have centeredScaled := centeredVariable.function at Ω, weightScale(Ω, c, X),
    expectation(Ω, P, weightScale(Ω, c, X)) with scaled
  have scaledCentered := weightScale.function at Ω, c,
    centeredVariable(Ω, X, expectation(Ω, P, X)) with centered
  have squareCentered := squareVariable.function with centered
  have squareScaledCentered := squareVariable.function with scaledCentered
  have squareCenteredScaled := squareVariable.function with centeredScaled
  have scaledSquare := weightScale.function at Ω, c * c,
    squareVariable(Ω, centeredVariable(Ω, X, expectation(Ω, P, X))) with
      squareCentered
  have agree : ∀ ω ∈ Ω, squareVariable(Ω, centeredVariable(Ω, weightScale(Ω, c,
      X),
      expectation(Ω, P, weightScale(Ω, c, X))))(ω) =
    weightScale(Ω, c * c, squareVariable(Ω, centeredVariable(Ω, X,
      expectation(Ω, P, X))))(ω)
    intros ω point
    have xReal := FunctionBetween.apply_closed with random, point
    have outer := squareVariable.value at Ω, centeredVariable(Ω, weightScale(Ω,
        c, X),
      expectation(Ω, P, weightScale(Ω, c, X))), ω with centeredScaled, point
    have inner := centeredVariable.value at Ω, weightScale(Ω, c, X),
      expectation(Ω, P, weightScale(Ω, c, X)), ω with scaled, point
    have scaledValue := weightScale.value at Ω, c, X, ω with random, point
    have rightOuter := weightScale.value at Ω, c * c,
      squareVariable(Ω, centeredVariable(Ω, X, expectation(Ω, P, X))), ω with
      squareCentered, point
    have rightSquare := squareVariable.value at Ω,
      centeredVariable(Ω, X, expectation(Ω, P, X)), ω with centered, point
    have rightCentered := centeredVariable.value at Ω, X, expectation(Ω, P, X),
        ω with
      random, point
    rw [outer, inner, scaledValue, scaledMean, rightOuter, rightSquare,
      rightCentered]
    ring
  have same := expectation.congr with space, squareCenteredScaled,
      scaledSquare, agree
  have pulled := expectation.scale at Ω, P,
    squareVariable(Ω, centeredVariable(Ω, X, expectation(Ω, P, X))), c * c with
    space, squareCentered
  rw [variance.value, variance.value, same, pulled]
  rfl

-- Shifting a random by a constant leaves its variance unchanged.
theorem variance.shift (Ω P X : Set) (c : Real)
  requires space : ProbabilitySpace(Ω, P)
  requires random : X : Ω → ℝ
  shows variance(Ω, P, weightAdd(Ω, X, constantFunction(Ω, ℝ,
      c))) = variance(Ω, P, X)
proof
  have mean := expectation.real with space, random
  have scalar : c ∈ ℝ by simp [Real]
  have constant := constantFunction.function at Ω, ℝ, c with scalar
  have shifted := weightAdd.function with random, constant
  have shiftedMean := expectation.add with space, random, constant
  have constantMean := expectation.constant at Ω, P, c with space
  rw constantMean at shiftedMean
  have centered := centeredVariable.function at Ω, X, expectation(Ω, P, X) with
      random
  have centeredShifted := centeredVariable.function at Ω,
    weightAdd(Ω, X, constantFunction(Ω, ℝ, c)),
    expectation(Ω, P, weightAdd(Ω, X, constantFunction(Ω, ℝ, c))) with shifted
  have squareCentered := squareVariable.function with centered
  have squareShifted := squareVariable.function with centeredShifted
  have agree : ∀ ω ∈ Ω, squareVariable(Ω, centeredVariable(Ω,
      weightAdd(Ω, X, constantFunction(Ω, ℝ, c)),
      expectation(Ω, P, weightAdd(Ω, X, constantFunction(Ω, ℝ, c)))))(ω) =
    squareVariable(Ω, centeredVariable(Ω, X, expectation(Ω, P, X)))(ω)
    intros ω point
    have xReal := FunctionBetween.apply_closed with random, point
    have outer := squareVariable.value at Ω, centeredVariable(Ω,
      weightAdd(Ω, X, constantFunction(Ω, ℝ, c)),
      expectation(Ω, P, weightAdd(Ω, X, constantFunction(Ω, ℝ, c)))), ω with
      centeredShifted, point
    have inner := centeredVariable.value at Ω, weightAdd(Ω, X,
        constantFunction(Ω, ℝ, c)),
      expectation(Ω, P, weightAdd(Ω, X, constantFunction(Ω, ℝ, c))), ω with
        shifted, point
    have sum := weightAdd.value at Ω, X, constantFunction(Ω, ℝ, c), ω with
      random, constant, point
    have constantValue := constantFunction.apply at Ω, ℝ, ω, c with point,
        scalar
    have rightSquare := squareVariable.value at Ω,
      centeredVariable(Ω, X, expectation(Ω, P, X)), ω with centered, point
    have rightCentered := centeredVariable.value at Ω, X, expectation(Ω, P, X),
        ω with
      random, point
    rw [outer, inner, sum, constantValue, shiftedMean, rightSquare,
      rightCentered]
    ring
  rw [variance.value, variance.value]
  exact expectation.congr with space, squareShifted, squareCentered, agree

theorem variance.constant (Ω P : Set) (c : Real)
  requires space : ProbabilitySpace(Ω, P)
  shows variance(Ω, P, constantFunction(Ω, ℝ, c)) = (0:R)
proof
  have scalar : c ∈ ℝ by simp [Real]
  have zeroReal := realZero.real
  have constant := constantFunction.function at Ω, ℝ, c with scalar
  have zeros := constantFunction.function at Ω, ℝ, (0:R) with zeroReal
  have constantMean := expectation.constant at Ω, P, c with space
  have centered := centeredVariable.function at Ω, constantFunction(Ω, ℝ, c),
    expectation(Ω, P, constantFunction(Ω, ℝ, c)) with constant
  have square := squareVariable.function with centered
  have agree : ∀ ω ∈ Ω, squareVariable(Ω, centeredVariable(Ω,
      constantFunction(Ω, ℝ, c),
      expectation(Ω, P, constantFunction(Ω, ℝ, c))))(ω) = constantFunction(Ω,
        ℝ, (0:R))(ω)
    intros ω point
    have outer := squareVariable.value at Ω, centeredVariable(Ω,
        constantFunction(Ω, ℝ, c),
      expectation(Ω, P, constantFunction(Ω, ℝ, c))), ω with centered, point
    have inner := centeredVariable.value at Ω, constantFunction(Ω, ℝ, c),
      expectation(Ω, P, constantFunction(Ω, ℝ, c)), ω with constant, point
    have constantValue := constantFunction.apply at Ω, ℝ, ω, c with point,
        scalar
    have zeroValue := constantFunction.apply at Ω, ℝ, ω, (0:R) with point,
        zeroReal
    rw [outer, inner, constantValue, constantMean, zeroValue]
    ring
  have same := expectation.congr with space, square, zeros, agree
  have zeroMean := expectation.constant at Ω, P, (0:R) with space
  rw [variance.value, same, zeroMean]
  rfl

-- Chebyshev's inequality in its squared form: the probability that the
-- squared deviation reaches a² is at most Var X / a².
theorem variance.chebyshev_square (Ω P X : Set) (a : Real)
  requires space : ProbabilitySpace(Ω, P)
  requires random : X : Ω → ℝ
  requires positive : (0:R) < a
  shows (prob(P, levelSet(Ω, squareVariable(Ω, centeredVariable(Ω, X,
      expectation(Ω, P, X))),
      a * a)):R) ≤ (variance(Ω, P, X):R) / (a * a)
proof
  have mean := expectation.real with space, random
  have centered := centeredVariable.function at Ω, X, expectation(Ω, P, X) with
      random
  have square := squareVariable.function with centered
  have nonnegative : ∀ ω ∈ Ω,
      (0:R) ≤ (squareVariable(Ω, centeredVariable(Ω, X, expectation(Ω, P,
        X)))(ω):R)
    intros ω point
    exact squareVariable.nonnegative with centered, point
  have squarePositive := realMulValue.positive at a, a with positive, positive
  have markov := expectation.markov at Ω, P,
    squareVariable(Ω, centeredVariable(Ω, X, expectation(Ω, P, X))), a * a with
    space, square, nonnegative, squarePositive
  rw variance.value
  exact markov

theorem abs.square_le (A B : Real)
  requires nonnegative : (0:R) ≤ A
  requires bound : A ≤ abs(B)
  shows A * A ≤ B * B
proof
  have absReal := abs.real at B
  have absNonnegative := abs.nonnegative at B
  have first := realLe.mul_right_of_nonnegative at A, abs(B), A with bound,
      nonnegative
  have second := realLe.mul_left_of_nonnegative at A, abs(B), abs(B) with
      bound, absNonnegative
  have squareNonnegative := realMulValue.square_nonnegative at B
  have product := abs.mul at B, B
  have same := abs.of_nonnegative at B * B with squareNonnegative
  have chain := realLe.transitive with first, second
  rw ← product at chain
  rwa same at chain

-- Chebyshev's inequality: the probability that X deviates from its mean by
-- at least a is at most Var X / a².
theorem variance.chebyshev (Ω P X : Set) (a : Real)
  requires space : ProbabilitySpace(Ω, P)
  requires random : X : Ω → ℝ
  requires positive : (0:R) < a
  shows (prob(P, levelSet(Ω, absVariable(Ω, centeredVariable(Ω, X,
      expectation(Ω, P, X))),
      a)):R) ≤ (variance(Ω, P, X):R) / (a * a)
proof
  have mean := expectation.real with space, random
  have centered := centeredVariable.function at Ω, X, expectation(Ω, P, X) with
      random
  have square := squareVariable.function with centered
  have absolute := absVariable.function with centered
  have nonnegative : (0:R) ≤ a by linarith
  have inner : levelSet(Ω, absVariable(Ω, centeredVariable(Ω, X, expectation(Ω,
      P, X))),
      a) ⊆ levelSet(Ω, squareVariable(Ω, centeredVariable(Ω, X, expectation(Ω,
        P, X))),
      a * a)
    intros ω member
    have data := levelSet.membership.forward with member
    have point := data.left
    have deviationReal := FunctionBetween.apply_closed with centered, point
    have absValue := absVariable.value at Ω,
      centeredVariable(Ω, X, expectation(Ω, P, X)), ω with centered, point
    have squareValue := squareVariable.value at Ω,
      centeredVariable(Ω, X, expectation(Ω, P, X)), ω with centered, point
    have bound := data.right
    rw absValue at bound
    have squared := abs.square_le at a, centeredVariable(Ω, X, expectation(Ω,
        P, X))(ω) with
      nonnegative, bound
    apply levelSet.membership.backward
    refine ⟨point, ?_⟩
    rw squareValue
    exact squared
  have eventSquare := levelSet.subset at Ω,
    squareVariable(Ω, centeredVariable(Ω, X, expectation(Ω, P, X))), a * a
  have monotone := prob.monotone with space, eventSquare, inner
  have chebyshev := variance.chebyshev_square with space, random, positive
  have eventAbs := levelSet.subset at Ω,
    absVariable(Ω, centeredVariable(Ω, X, expectation(Ω, P, X))), a
  have realAbs := prob.real with space, eventAbs
  have realSquare := prob.real with space, eventSquare
  have varianceReal := variance.real with space, random
  have squareReal := realMulValue.real at a, a
  have quotientReal := realDiv.real at variance(Ω, P, X), a * a
  exact realLe.transitive with monotone, chebyshev