-- The ballot problem as a probability: under the uniform distribution on
-- the orderings of a votes for A and b < a votes for B, A stays strictly
-- ahead throughout the count with probability (a − b)/(a + b).
-- 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.Combinatorics.BallotProbability
import Ballot
import Probability.Uniform

-- Bertrand's ballot theorem, as a probability.
theorem ballot.probability (a b : Nat)
  requires lead : b < a
  shows prob(uniformWeight(orderings(a, b)), ballotSequences(a, b)) =
    (realNat(a - b):R) / (realNat(a + b):R)
proof
  have finite := orderings.finite at a, b
  have event := ballotSequences.subset_orderings at a, b
  have classical := prob.uniform with finite, event
  have count := ballot.count with lead
  have le : a ≤ a + b by arith
  have positive := choose.positive with le
  have nonzeroChoose : (a + b) choose a ≠ 0
    intro zero
    rw zero at positive
    arith
  have nonzeroTotal : a + b ≠ 0
    intro zero
    arith
  have realChoose := realNat.ne_zero with nonzeroChoose
  have realTotal := realNat.ne_zero with nonzeroTotal
  have cross : (realNat(a - b):R) * (realNat((a + b) choose a):R) =
      (realNat(a + b):R) * (realNat(cardinality(ballotSequences(a, b))):R)
    rw [← realNat.mul, ← realNat.mul, count]
    rfl
  rw [classical, realSize.value, realSize.value, orderings.count]
  exact realDiv.eq_of_cross with realChoose, realTotal, cross