-- Point masses: the probability space concentrated at one outcome
-- 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.Dirac
import Variance

-- The weight that puts all mass on the single outcome ω₀.
def diracWeight (Ω ω₀ : Set) : Set := realIndicator(Ω, {ω₀})

theorem diracWeight.function (Ω ω₀ : Set)
  shows diracWeight(Ω, ω₀) : Ω → ℝ
proof
  unfold diracWeight
  exact realIndicator.function at Ω, {ω₀}

theorem diracWeight.at_point (Ω ω₀ : Set)
  requires point : ω₀ ∈ Ω
  shows diracWeight(Ω, ω₀)(ω₀) = (1:R)
proof
  have self := singleton.self at ω₀
  unfold diracWeight
  exact realIndicator.inside at Ω, {ω₀}, ω₀ with point, self

theorem diracWeight.elsewhere (Ω ω₀ ω : Set)
  requires point : ω ∈ Ω
  requires different : ω ≠ ω₀
  shows diracWeight(Ω, ω₀)(ω) = (0:R)
proof
  have outside : ¬(ω ∈ {ω₀}) by simp [singleton.membership]
  unfold diracWeight
  exact realIndicator.outside at Ω, {ω₀}, ω with point, outside

theorem diracWeight.nonnegative (Ω ω₀ : Set)
  shows ∀ ω ∈ Ω, (0:R) ≤ (diracWeight(Ω, ω₀)(ω):R)
proof
  intros ω point
  unfold diracWeight
  exact realIndicator.nonnegative at Ω, {ω₀}, ω with point

-- Summing the point mass over an event counts whether ω₀ lies in it.
theorem diracWeight.sum_of_member (Ω ω₀ A : Set)
  requires finite : Finite(A)
  requires event : A ⊆ Ω
  requires member : ω₀ ∈ A
  shows realSetSum(A, diracWeight(Ω, ω₀)) = (1:R)
proof
  have counted := realSetSum.indicator at A, {ω₀}, Ω with finite, event
  have meet : A ∩ {ω₀} = {ω₀}
    ext y
    simp
  have one := cardinality.singleton at ω₀
  simp [diracWeight]
  exact realNat.one

theorem diracWeight.sum_of_absent (Ω ω₀ A : Set)
  requires finite : Finite(A)
  requires event : A ⊆ Ω
  requires absent : ¬(ω₀ ∈ A)
  shows realSetSum(A, diracWeight(Ω, ω₀)) = (0:R)
proof
  have counted := realSetSum.indicator at A, {ω₀}, Ω with finite, event
  have meet : A ∩ {ω₀} = ∅
    ext y
    simp
  have zero := cardinality.eq_of_has with HasCardinality.empty_set
  simp [diracWeight]

theorem diracWeight.total (Ω ω₀ : Set)
  requires finite : Finite(Ω)
  requires point : ω₀ ∈ Ω
  shows realSetSum(Ω, diracWeight(Ω, ω₀)) = (1:R)
proof
  have whole := Subset.reflexive at Ω
  exact diracWeight.sum_of_member with finite, whole, point

theorem diracWeight.space (Ω ω₀ : Set)
  requires finite : Finite(Ω)
  requires point : ω₀ ∈ Ω
  shows ProbabilitySpace(Ω, diracWeight(Ω, ω₀))
proof
  have function := diracWeight.function at Ω, ω₀
  have nonnegative := diracWeight.nonnegative at Ω, ω₀
  have total := diracWeight.total with finite, point
  exact ProbabilitySpace.intro with finite, function, nonnegative, total

theorem prob.dirac_of_member (Ω ω₀ A : Set)
  requires finite : Finite(Ω)
  requires event : A ⊆ Ω
  requires member : ω₀ ∈ A
  shows prob(diracWeight(Ω, ω₀), A) = (1:R)
proof
  have finiteA := Finite.subset with finite, event
  rw prob.value
  exact diracWeight.sum_of_member with finiteA, event, member

theorem prob.dirac_of_absent (Ω ω₀ A : Set)
  requires finite : Finite(Ω)
  requires event : A ⊆ Ω
  requires absent : ¬(ω₀ ∈ A)
  shows prob(diracWeight(Ω, ω₀), A) = (0:R)
proof
  have finiteA := Finite.subset with finite, event
  rw prob.value
  exact diracWeight.sum_of_absent with finiteA, event, absent

-- Under a point mass, every random variable is constant in expectation.
theorem expectation.dirac (Ω ω₀ X : Set)
  requires finite : Finite(Ω)
  requires point : ω₀ ∈ Ω
  requires random : X : Ω → ℝ
  shows expectation(Ω, diracWeight(Ω, ω₀), X) = X(ω₀)
proof
  have whole := Subset.reflexive at Ω
  have function := diracWeight.function at Ω, ω₀
  have product := weightMul.function with function, random
  have masked := maskedWeight.function at Ω, {ω₀}, X with random
  have agree : ∀ ω ∈ Ω, weightMul(Ω, diracWeight(Ω, ω₀), X)(ω) =
      maskedWeight(Ω, {ω₀}, X)(ω)
    intros ω inside
    have xReal := FunctionBetween.apply_closed with random, inside
    have natural := indicator.natural at {ω₀}, ω
    have maskReal := realNat.real at indicator({ω₀}, ω)
    have left := weightMul.value at Ω, diracWeight(Ω, ω₀), X, ω with function,
        random, inside
    have inner : diracWeight(Ω, ω₀)(ω) = realNat(indicator({ω₀}, ω))
      unfold diracWeight
      exact realIndicator.value at Ω, {ω₀}, ω with inside
    have right := maskedWeight.value at Ω, {ω₀}, X, ω with random, inside
    rw [left, inner, right]
    ring
  have same := realSetSum.congr with product, masked, whole, agree
  have counted := realSetSum.masked_intersection at Ω, {ω₀}, Ω, X with finite,
      random, whole
  have meet : Ω ∩ {ω₀} = {ω₀}
    ext y
    simp
  have single := realSetSum.singleton with random, point
  rw [expectation.value, same, counted, meet, single]
  rfl

theorem variance.dirac (Ω ω₀ X : Set)
  requires finite : Finite(Ω)
  requires point : ω₀ ∈ Ω
  requires random : X : Ω → ℝ
  shows variance(Ω, diracWeight(Ω, ω₀), X) = (0:R)
proof
  have mean := expectation.dirac with finite, point, random
  have xReal := FunctionBetween.apply_closed with random, point
  have centered := centeredVariable.function at Ω, X, expectation(Ω,
      diracWeight(Ω, ω₀), X)
    with random
  have square := squareVariable.function with centered
  have squareMean := expectation.dirac at Ω, ω₀,
    squareVariable(Ω, centeredVariable(Ω, X, expectation(Ω, diracWeight(Ω, ω₀),
      X))) with
    finite, point, square
  have squareValue := squareVariable.value at Ω,
    centeredVariable(Ω, X, expectation(Ω, diracWeight(Ω, ω₀), X)), ω₀ with
      centered, point
  have centeredValue := centeredVariable.value at Ω, X, expectation(Ω,
      diracWeight(Ω, ω₀), X),
    ω₀ with random, point
  rw [variance.value, squareMean, squareValue, centeredValue, mean]
  ring

theorem diracWeight.prob_point (Ω ω₀ : Set)
  requires finite : Finite(Ω)
  requires point : ω₀ ∈ Ω
  shows prob(diracWeight(Ω, ω₀), {ω₀}) = (1:R)
proof
  have event : {ω₀} ⊆ Ω by simp
  have self := singleton.self at ω₀
  exact prob.dirac_of_member with finite, event, self