-- The uniform distribution on a nonempty finite set
-- 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.Uniform
import Space

-- The number of points of a finite set, as a real.
def realSize (Ω : Set) : Set := realNat(cardinality(Ω))

theorem realSize.value (Ω : Set)
  shows realSize(Ω) = realNat(cardinality(Ω))
proof
  simp [realSize]

theorem realSize.real (Ω : Set)
  requires finite : Finite(Ω)
  shows Real(realSize(Ω))
proof
  have natural := cardinality.nat with finite
  rw realSize.value
  exact realNat.real at cardinality(Ω)

theorem realSize.nonnegative (Ω : Set)
  requires finite : Finite(Ω)
  shows (0:R) ≤ (realSize(Ω):R)
proof
  have natural := cardinality.nat with finite
  rw realSize.value
  exact realNat.zero_le at cardinality(Ω)

theorem cardinality.ne_zero_of_nonempty (Ω : Set)
  requires finite : Finite(Ω)
  requires inhabited : Nonempty(Ω)
  shows cardinality(Ω) ≠ 0
proof
  intro zero
  have correct := cardinality.correct with finite
  rw zero at correct
  have empty := HasCardinality.empty_of_zero with correct
  have different := Nonempty.ne_empty with inhabited
  exact different with empty

theorem realSize.positive (Ω : Set)
  requires finite : Finite(Ω)
  requires inhabited : Nonempty(Ω)
  shows (0:R) < (realSize(Ω):R)
proof
  have natural := cardinality.nat with finite
  have nonzero := cardinality.ne_zero_of_nonempty with finite, inhabited
  rw realSize.value
  exact realNat.positive at cardinality(Ω) with nonzero

theorem realSize.nonzero (Ω : Set)
  requires finite : Finite(Ω)
  requires inhabited : Nonempty(Ω)
  shows realSize(Ω) ≠ (0:R)
proof
  have sizeReal := realSize.real with finite
  have positive := realSize.positive with finite, inhabited
  intro zero
  rw zero at positive
  exact realLess.not_self with positive

theorem realSize.monotone (A B : Set)
  requires finite : Finite(B)
  requires subset : A ⊆ B
  shows (realSize(A):R) ≤ (realSize(B):R)
proof
  have finiteA := Finite.subset with finite, subset
  have naturalA := cardinality.nat with finiteA
  have naturalB := cardinality.nat with finite
  have le := cardinality.monotone with finite, subset
  rw [realSize.value, realSize.value]
  exact realNat.le_of_natural_le at cardinality(A), cardinality(B) with le

-- The uniform weight gives every point the mass 1 / |Ω|.
def uniformWeight (Ω : Set) : Set :=
  constantFunction(Ω, ℝ, realInvValue(realSize(Ω)))

theorem uniformWeight.mass_real (Ω : Set)
  requires finite : Finite(Ω)
  shows Real(realInvValue(realSize(Ω)))
proof
  have sizeReal := realSize.real with finite
  exact realInvValue.real at realSize(Ω)

theorem uniformWeight.function (Ω : Set)
  requires finite : Finite(Ω)
  shows uniformWeight(Ω) : Ω → ℝ
proof
  have massReal := uniformWeight.mass_real with finite
  unfold uniformWeight
  exact constantFunction.function at Ω, ℝ, realInvValue(realSize(Ω)) with
    massReal

theorem uniformWeight.value (Ω ω : Set)
  requires finite : Finite(Ω)
  requires point : ω ∈ Ω
  shows uniformWeight(Ω)(ω) = realInvValue(realSize(Ω))
proof
  have massReal := uniformWeight.mass_real with finite
  unfold uniformWeight
  exact constantFunction.apply at Ω, ℝ, ω, realInvValue(realSize(Ω)) with
    point, massReal

theorem uniformWeight.mass_positive (Ω : Set)
  requires finite : Finite(Ω)
  requires inhabited : Nonempty(Ω)
  shows (0:R) < (realInvValue(realSize(Ω)):R)
proof
  have sizeReal := realSize.real with finite
  have positive := realSize.positive with finite, inhabited
  exact realInvValue.positive at realSize(Ω) with positive

theorem uniformWeight.nonnegative (Ω : Set)
  requires finite : Finite(Ω)
  requires inhabited : Nonempty(Ω)
  shows ∀ ω ∈ Ω, (0:R) ≤ (uniformWeight(Ω)(ω):R)
proof
  intros ω point
  have value := uniformWeight.value with finite, point
  have massReal := uniformWeight.mass_real with finite
  have positive := uniformWeight.mass_positive with finite, inhabited
  rw value
  linarith

-- Summing the uniform weight over an event gives |A| / |Ω|.
theorem uniformWeight.sum (Ω A : Set)
  requires finite : Finite(Ω)
  requires event : A ⊆ Ω
  shows realSetSum(A,
      uniformWeight(Ω)) = (realSize(A):R) * (realInvValue(realSize(Ω)):R)
proof
  have massReal := uniformWeight.mass_real with finite
  have finiteA := Finite.subset with finite, event
  have constant := realSetSum.constant at A, Ω, realInvValue(realSize(Ω)) with
    finiteA, event
  have size := realSize.value at A
  unfold uniformWeight
  cc

theorem uniformWeight.total (Ω : Set)
  requires finite : Finite(Ω)
  requires inhabited : Nonempty(Ω)
  shows realSetSum(Ω, uniformWeight(Ω)) = (1:R)
proof
  have whole := Subset.reflexive at Ω
  have sum := uniformWeight.sum with finite, whole
  have sizeReal := realSize.real with finite
  have nonzero := realSize.nonzero with finite, inhabited
  have inverse := realInvValue.mul_right at realSize(Ω) with nonzero
  rw sum
  exact inverse

theorem uniformWeight.space (Ω : Set)
  requires finite : Finite(Ω)
  requires inhabited : Nonempty(Ω)
  shows ProbabilitySpace(Ω, uniformWeight(Ω))
proof
  have function := uniformWeight.function with finite
  have nonnegative := uniformWeight.nonnegative with finite, inhabited
  have total := uniformWeight.total with finite, inhabited
  exact ProbabilitySpace.intro with finite, function, nonnegative, total

-- Classical probability: favourable cases over all cases.
theorem prob.uniform (Ω A : Set)
  requires finite : Finite(Ω)
  requires event : A ⊆ Ω
  shows prob(uniformWeight(Ω), A) = (realSize(A):R) / (realSize(Ω):R)
proof
  have finiteA := Finite.subset with finite, event
  have sizeReal := realSize.real with finite
  have sizeAReal := realSize.real with finiteA
  have sum := uniformWeight.sum with finite, event
  have division := realDiv.eq_mul_inv at realSize(A), realSize(Ω)
  rw [prob.value, sum, division]
  rfl

theorem prob.uniform_point (Ω ω : Set)
  requires finite : Finite(Ω)
  requires point : ω ∈ Ω
  shows prob(uniformWeight(Ω), {ω}) = (1:R) / (realSize(Ω):R)
proof
  have event : {ω} ⊆ Ω by simp
  have classical := prob.uniform with finite, event
  have one : realSize({ω}) = (1:R)
    rw [realSize.value, cardinality.singleton]
    exact realNat.one
  rwa one at classical

theorem prob.uniform_whole (Ω : Set)
  requires finite : Finite(Ω)
  requires inhabited : Nonempty(Ω)
  shows prob(uniformWeight(Ω), Ω) = (1:R)
proof
  have space := uniformWeight.space with finite, inhabited
  exact prob.whole with space

theorem prob.uniform_complement (Ω A : Set)
  requires finite : Finite(Ω)
  requires inhabited : Nonempty(Ω)
  requires event : A ⊆ Ω
  shows prob(uniformWeight(Ω),
      Ω ∖ A) = (1:R) - (realSize(A):R) / (realSize(Ω):R)
proof
  have space := uniformWeight.space with finite, inhabited
  have complement := prob.complement with space, event
  have classical := prob.uniform with finite, event
  rwa classical at complement

-- Events of the same size are equally likely.
theorem prob.uniform_congr (Ω A B : Set)
  requires finite : Finite(Ω)
  requires eventA : A ⊆ Ω
  requires eventB : B ⊆ Ω
  requires same : cardinality(A) = cardinality(B)
  shows prob(uniformWeight(Ω), A) = prob(uniformWeight(Ω), B)
proof
  have classicalA := prob.uniform with finite, eventA
  have classicalB := prob.uniform with finite, eventB
  have sizes : realSize(A) = realSize(B)
    rw [realSize.value, realSize.value, same]
    rfl
  rw [classicalA, classicalB, sizes]
  rfl

theorem prob.uniform_le_of_size (Ω A B : Set)
  requires finite : Finite(Ω)
  requires eventA : A ⊆ Ω
  requires eventB : B ⊆ Ω
  requires smaller : LessEqual(cardinality(A), cardinality(B))
  shows (prob(uniformWeight(Ω), A):R) ≤ (prob(uniformWeight(Ω), B):R)
proof
  have finiteA := Finite.subset with finite, eventA
  have finiteB := Finite.subset with finite, eventB
  have naturalA := cardinality.nat with finiteA
  have naturalB := cardinality.nat with finiteB
  have sizeReal := realSize.real with finite
  have sizeAReal := realSize.real with finiteA
  have sizeBReal := realSize.real with finiteB
  have massReal := uniformWeight.mass_real with finite
  have classicalA := prob.uniform with finite, eventA
  have classicalB := prob.uniform with finite, eventB
  have divisionA := realDiv.eq_mul_inv at realSize(A), realSize(Ω)
  have divisionB := realDiv.eq_mul_inv at realSize(B), realSize(Ω)
  have inverseNonnegative : (0:R) ≤ (realInvValue(realSize(Ω)):R)
    have alternatives : Nonempty(Ω) ∨ ¬(Nonempty(Ω)) by tauto
    cases alternatives with inhabited | empty
    . have positive := uniformWeight.mass_positive with finite, inhabited
      linarith
    . have isEmpty : Ω = ∅
        have subset : Ω ⊆ ∅
          intros x hx
          exfalso
          apply empty
          exact ⟨x, hx⟩
        exact Subset.antisymmetric with subset, (empty.subset at Ω)
      have emptyCount := cardinality.eq_of_has with HasCardinality.empty_set
      have zero : realSize(Ω) = (0:R)
        rw [realSize.value, isEmpty, emptyCount]
        exact realNat.zero
      rw [zero, realInvValue.zero]
      exact realLe.reflexive with realZero.real
  have sizeLe : (realSize(A):R) ≤ (realSize(B):R)
    rw [realSize.value, realSize.value]
    exact realNat.le_of_natural_le at cardinality(A), cardinality(B) with smaller
  have scaled := realLe.mul_right_of_nonnegative at realSize(A), realSize(B),
    realInvValue(realSize(Ω)) with sizeLe, inverseNonnegative
  rw [classicalA, classicalB, divisionA, divisionB]
  exact scaled