-- Bernoulli trials, a fair coin, and a fair die
-- 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.Bernoulli
import Variance
import Uniform

-- The two-point sample space of a single trial.
def twoPoint : Set := {(0:N), (1:N)}

theorem twoPoint.membership (x : Set)
  shows x ∈ twoPoint ↔ (x = (0:N) ∨ x = (1:N))
proof
  simp [twoPoint, pair.membership]

theorem twoPoint.zero
  shows (0:N) ∈ twoPoint
proof
  simp [twoPoint.membership]

theorem twoPoint.one
  shows (1:N) ∈ twoPoint
proof
  simp [twoPoint.membership]

theorem twoPoint.as_insert
  shows twoPoint = insert((0:N), {(1:N)})
proof
  ext y
  simp [twoPoint.membership]

theorem twoPoint.finite
  shows Finite(twoPoint)
proof
  have single := Finite.singleton at (1:N)
  have inserted := Finite.insert at {(1:N)}, (0:N) with single
  rw twoPoint.as_insert
  exact inserted

theorem twoPoint.nonempty
  shows Nonempty(twoPoint)
proof
  have member := twoPoint.zero
  exact ⟨(0:N), member⟩

theorem twoPoint.cardinality
  shows cardinality(twoPoint) = 2
proof
  unfold twoPoint
  decide

theorem twoPoint.zero_fresh
  shows ¬((0:N) ∈ {(1:N)})
proof
  decide

theorem twoPoint.one_fresh
  shows ¬((1:N) ∈ {(0:N)})
proof
  decide

-- The Bernoulli weight with success probability p: mass p at 1 and 1 - p
-- at 0.
def bernoulliWeight (p : Real) : Set :=
  weightAdd(twoPoint, weightScale(twoPoint, p, realIndicator(twoPoint,
    {(1:N)})),
    weightScale(twoPoint, (1:R) - p, realIndicator(twoPoint, {(0:N)})))

theorem bernoulliWeight.function (p : Real)
  shows bernoulliWeight(p) : twoPoint → ℝ
proof
  have failure := realSubValue.real at (1:R), p
  have successes := realIndicator.function at twoPoint, {(1:N)}
  have failures := realIndicator.function at twoPoint, {(0:N)}
  have scaledSuccess := weightScale.function at twoPoint, p,
    realIndicator(twoPoint, {(1:N)}) with successes
  have scaledFailure := weightScale.function at twoPoint, (1:R) - p,
    realIndicator(twoPoint, {(0:N)}) with failures
  unfold bernoulliWeight
  exact weightAdd.function with scaledSuccess, scaledFailure

theorem bernoulliWeight.at_one (p : Real)
  shows bernoulliWeight(p)((1:N)) = p
proof
  have failure := realSubValue.real at (1:R), p
  have successes := realIndicator.function at twoPoint, {(1:N)}
  have failures := realIndicator.function at twoPoint, {(0:N)}
  have scaledSuccess := weightScale.function at twoPoint, p,
    realIndicator(twoPoint, {(1:N)}) with successes
  have scaledFailure := weightScale.function at twoPoint, (1:R) - p,
    realIndicator(twoPoint, {(0:N)}) with failures
  have point := twoPoint.one
  have sum := weightAdd.value at twoPoint, weightScale(twoPoint, p,
      realIndicator(twoPoint, {(1:N)})),
    weightScale(twoPoint, (1:R) - p, realIndicator(twoPoint, {(0:N)})),
      (1:N) with
    scaledSuccess, scaledFailure, point
  have first := weightScale.value at twoPoint, p, realIndicator(twoPoint,
      {(1:N)}), (1:N) with
    successes, point
  have second := weightScale.value at twoPoint, (1:R) - p,
      realIndicator(twoPoint, {(0:N)}),
    (1:N) with failures, point
  have self := singleton.self at (1:N)
  have fresh := twoPoint.one_fresh
  have inside := realIndicator.inside at twoPoint, {(1:N)}, (1:N) with point,
      self
  have outside := realIndicator.outside at twoPoint, {(0:N)}, (1:N) with point,
      fresh
  simp [bernoulliWeight]
  ring

theorem bernoulliWeight.at_zero (p : Real)
  shows bernoulliWeight(p)((0:N)) = (1:R) - p
proof
  have failure := realSubValue.real at (1:R), p
  have successes := realIndicator.function at twoPoint, {(1:N)}
  have failures := realIndicator.function at twoPoint, {(0:N)}
  have scaledSuccess := weightScale.function at twoPoint, p,
    realIndicator(twoPoint, {(1:N)}) with successes
  have scaledFailure := weightScale.function at twoPoint, (1:R) - p,
    realIndicator(twoPoint, {(0:N)}) with failures
  have point := twoPoint.zero
  have sum := weightAdd.value at twoPoint, weightScale(twoPoint, p,
      realIndicator(twoPoint, {(1:N)})),
    weightScale(twoPoint, (1:R) - p, realIndicator(twoPoint, {(0:N)})),
      (0:N) with
    scaledSuccess, scaledFailure, point
  have first := weightScale.value at twoPoint, p, realIndicator(twoPoint,
      {(1:N)}), (0:N) with
    successes, point
  have second := weightScale.value at twoPoint, (1:R) - p,
      realIndicator(twoPoint, {(0:N)}),
    (0:N) with failures, point
  have self := singleton.self at (0:N)
  have fresh := twoPoint.zero_fresh
  have outside := realIndicator.outside at twoPoint, {(1:N)}, (0:N) with point,
      fresh
  have inside := realIndicator.inside at twoPoint, {(0:N)}, (0:N) with point,
      self
  simp [bernoulliWeight]
  ring

theorem bernoulliWeight.nonnegative (p : Real)
  requires lower : (0:R) ≤ p
  requires upper : p ≤ (1:R)
  shows ∀ ω ∈ twoPoint, (0:R) ≤ (bernoulliWeight(p)(ω):R)
proof
  intros ω point
  have alternatives := twoPoint.membership.forward with point
  cases alternatives with zero | one
  . have value := bernoulliWeight.at_zero at p
    rw [zero, value]
    linarith
  . have value := bernoulliWeight.at_one at p
    rw [one, value]
    exact lower

theorem bernoulliWeight.total (p : Real)
  shows realSetSum(twoPoint, bernoulliWeight(p)) = (1:R)
proof
  have function := bernoulliWeight.function at p
  have single := Finite.singleton at (1:N)
  have subset : {(1:N)} ⊆ twoPoint
    intros y member
    have equal := singleton.membership.forward with member
    rw equal
    exact twoPoint.one
  have zeroIn := twoPoint.zero
  have oneIn := twoPoint.one
  have fresh := twoPoint.zero_fresh
  have inserted := realSetSum.insert with single, function, subset, zeroIn,
      fresh
  have singleton := realSetSum.singleton with function, oneIn
  have atZero := bernoulliWeight.at_zero at p
  have atOne := bernoulliWeight.at_one at p
  rw twoPoint.as_insert
  rw [inserted, singleton, atZero, atOne]
  ring

theorem bernoulliWeight.space (p : Real)
  requires lower : (0:R) ≤ p
  requires upper : p ≤ (1:R)
  shows ProbabilitySpace(twoPoint, bernoulliWeight(p))
proof
  have finite := twoPoint.finite
  have function := bernoulliWeight.function at p
  have nonnegative := bernoulliWeight.nonnegative with lower, upper
  have total := bernoulliWeight.total at p
  exact ProbabilitySpace.intro with finite, function, nonnegative, total

theorem bernoulliWeight.prob_one (p : Real)
  requires lower : (0:R) ≤ p
  requires upper : p ≤ (1:R)
  shows prob(bernoulliWeight(p), {(1:N)}) = p
proof
  have space := bernoulliWeight.space with lower, upper
  have point := twoPoint.one
  have singleton := prob.singleton with space, point
  rw singleton
  exact bernoulliWeight.at_one at p

theorem bernoulliWeight.prob_zero (p : Real)
  requires lower : (0:R) ≤ p
  requires upper : p ≤ (1:R)
  shows prob(bernoulliWeight(p), {(0:N)}) = (1:R) - p
proof
  have space := bernoulliWeight.space with lower, upper
  have point := twoPoint.zero
  have singleton := prob.singleton with space, point
  rw singleton
  exact bernoulliWeight.at_zero at p

-- The outcome of the trial as a random variable: one on success.
def bernoulliVariable : Set := realIndicator(twoPoint, {(1:N)})

theorem bernoulliVariable.function
  shows bernoulliVariable : twoPoint → ℝ
proof
  unfold bernoulliVariable
  exact realIndicator.function at twoPoint, {(1:N)}

theorem bernoulliVariable.expectation (p : Real)
  requires lower : (0:R) ≤ p
  requires upper : p ≤ (1:R)
  shows expectation(twoPoint, bernoulliWeight(p), bernoulliVariable) = p
proof
  have space := bernoulliWeight.space with lower, upper
  have event : {(1:N)} ⊆ twoPoint
    intros y member
    have equal := singleton.membership.forward with member
    rw equal
    exact twoPoint.one
  have indicator := expectation.indicator with space, event
  have success := bernoulliWeight.prob_one with lower, upper
  unfold bernoulliVariable
  cc

-- An indicator is its own square.
theorem bernoulliVariable.square_expectation (p : Real)
  requires lower : (0:R) ≤ p
  requires upper : p ≤ (1:R)
  shows expectation(twoPoint, bernoulliWeight(p), squareVariable(twoPoint,
      bernoulliVariable)) = p
proof
  have space := bernoulliWeight.space with lower, upper
  have random := bernoulliVariable.function
  have square := squareVariable.function with random
  have agree : ∀ ω ∈ twoPoint, squareVariable(twoPoint, bernoulliVariable)(ω) =
      bernoulliVariable(ω)
    intros ω point
    have value := squareVariable.value at twoPoint, bernoulliVariable, ω with
        random, point
    have alternatives : ω ∈ {(1:N)} ∨ ¬(ω ∈ {(1:N)}) by tauto
    cases alternatives with inside | outside
    . have one : bernoulliVariable(ω) = (1:R)
        unfold bernoulliVariable
        exact realIndicator.inside at twoPoint, {(1:N)}, ω with point, inside
      rw [value, one]
      ring
    . have zero : bernoulliVariable(ω) = (0:R)
        unfold bernoulliVariable
        exact realIndicator.outside at twoPoint, {(1:N)}, ω with point, outside
      rw [value, zero]
      ring
  have same := expectation.congr with space, square, random, agree
  rw same
  exact bernoulliVariable.expectation with lower, upper

theorem bernoulliVariable.variance (p : Real)
  requires lower : (0:R) ≤ p
  requires upper : p ≤ (1:R)
  shows variance(twoPoint, bernoulliWeight(p),
      bernoulliVariable) = p * ((1:R) - p)
proof
  have space := bernoulliWeight.space with lower, upper
  have random := bernoulliVariable.function
  have formula := variance.formula with space, random
  have mean := bernoulliVariable.expectation with lower, upper
  have squareMean := bernoulliVariable.square_expectation with lower, upper
  rw [formula, mean, squareMean]
  ring

-- A fair coin is the uniform distribution on two points.
def fairCoin : Set := uniformWeight(twoPoint)

theorem fairCoin.space
  shows ProbabilitySpace(twoPoint, fairCoin)
proof
  have finite := twoPoint.finite
  have inhabited := twoPoint.nonempty
  unfold fairCoin
  exact uniformWeight.space with finite, inhabited

theorem twoPoint.size
  shows realSize(twoPoint) = (2:R)
proof
  have numeral := realNumeral.nat at 2
  rw [realSize.value, twoPoint.cardinality, numeral]
  rfl

theorem fairCoin.heads
  shows prob(fairCoin, {(1:N)}) = (1:R) / (2:R)
proof
  have finite := twoPoint.finite
  have point := twoPoint.one
  have classical := prob.uniform_point with finite, point
  have size := twoPoint.size
  unfold fairCoin
  cc

theorem fairCoin.tails
  shows prob(fairCoin, {(0:N)}) = (1:R) / (2:R)
proof
  have finite := twoPoint.finite
  have point := twoPoint.zero
  have classical := prob.uniform_point with finite, point
  have size := twoPoint.size
  unfold fairCoin
  cc

theorem fairCoin.expectation
  shows expectation(twoPoint, fairCoin, bernoulliVariable) = (1:R) / (2:R)
proof
  have space := fairCoin.space
  have event : {(1:N)} ⊆ twoPoint
    intros y member
    have equal := singleton.membership.forward with member
    rw equal
    exact twoPoint.one
  have indicator := expectation.indicator with space, event
  have heads := fairCoin.heads
  unfold bernoulliVariable
  cc

-- A fair die is the uniform distribution on six faces.
def dieFaces : Set := {(1:N), (2:N), (3:N), (4:N), (5:N), (6:N)}

theorem dieFaces.bounded
  shows dieFaces ⊆ (7:N)
proof
  unfold dieFaces
  decide

theorem dieFaces.finite
  shows Finite(dieFaces)
proof
  have seven := Finite.natural at (7:N)
  have bounded := dieFaces.bounded
  exact Finite.subset with seven, bounded

theorem dieFaces.one
  shows (1:N) ∈ dieFaces
proof
  unfold dieFaces
  decide

theorem dieFaces.nonempty
  shows Nonempty(dieFaces)
proof
  have member := dieFaces.one
  exact ⟨(1:N), member⟩

theorem dieFaces.cardinality
  shows cardinality(dieFaces) = 6
proof
  unfold dieFaces
  decide

theorem dieFaces.size
  shows realSize(dieFaces) = (6:R)
proof
  have numeral := realNumeral.nat at 6
  rw [realSize.value, dieFaces.cardinality, numeral]
  rfl

def fairDie : Set := uniformWeight(dieFaces)

theorem fairDie.space
  shows ProbabilitySpace(dieFaces, fairDie)
proof
  have finite := dieFaces.finite
  have inhabited := dieFaces.nonempty
  unfold fairDie
  exact uniformWeight.space with finite, inhabited

theorem fairDie.six
  shows prob(fairDie, {(6:N)}) = (1:R) / (6:R)
proof
  have finite := dieFaces.finite
  have point : (6:N) ∈ dieFaces
    unfold dieFaces
    decide
  have classical := prob.uniform_point with finite, point
  have size := dieFaces.size
  unfold fairDie
  cc

def evenFaces : Set := {(2:N), (4:N), (6:N)}

theorem evenFaces.event
  shows evenFaces ⊆ dieFaces
proof
  unfold evenFaces
  unfold dieFaces
  decide

theorem evenFaces.cardinality
  shows cardinality(evenFaces) = 3
proof
  unfold evenFaces
  decide

theorem evenFaces.size
  shows realSize(evenFaces) = (3:R)
proof
  have numeral := realNumeral.nat at 3
  rw [realSize.value, evenFaces.cardinality, numeral]
  rfl

theorem fairDie.even
  shows prob(fairDie, evenFaces) = (3:R) / (6:R)
proof
  have finite := dieFaces.finite
  have event := evenFaces.event
  have classical := prob.uniform with finite, event
  have sizeEven := evenFaces.size
  have sizeAll := dieFaces.size
  unfold fairDie
  cc

theorem realDiv.three_sixths
  shows (3:R) / (6:R) = (1:R) / (2:R)
proof
  have twoReal := realNumeral.real at 2
  have threeReal := realNumeral.real at 3
  have sixReal := realNumeral.real at 6
  have twoPositive : (0:R) < (2:R) by linarith
  have sixPositive : (0:R) < (6:R) by linarith
  have twoNonzero : (2:R) ≠ (0:R)
    intro zero
    rw zero at twoPositive
    exact realLess.not_self with twoPositive
  have sixNonzero : (6:R) ≠ (0:R)
    intro zero
    rw zero at sixPositive
    exact realLess.not_self with sixPositive
  have inverseSixReal := realInvValue.real at (6:R)
  have inverseTwoReal := realInvValue.real at (2:R)
  have sixLaw := realInvValue.mul_right at (6:R) with sixNonzero
  have divisionThree := realDiv.eq_mul_inv at (3:R), (6:R)
  have divisionOne := realDiv.eq_mul_inv at (1:R), (2:R)
  have product : (2:R) * ((3:R) * (realInvValue((6:R)):R)) = (1:R)
    calc
      (2:R) * ((3:R) * (realInvValue((6:R)):R)) =
        (6:R) * (realInvValue((6:R)):R) := by ring
      _ = (1:R) := sixLaw
  have inverse := realInvValue.eq_of_mul_right at (2:R),
      (3:R) * (realInvValue((6:R)):R) with
    twoNonzero, product
  have oneLeft := realMulValue.one_left at realInvValue((2:R))
  rw [divisionThree, divisionOne, inverse, oneLeft]
  rfl

theorem fairDie.even_half
  shows prob(fairDie, evenFaces) = (1:R) / (2:R)
proof
  have even := fairDie.even
  have same := realDiv.three_sixths
  rw [even, same]
  rfl