-- Bertrand's ballot problem: with a votes for A and b < a votes for B, the
-- orderings in which A stays strictly ahead throughout the count number
-- (a − b)/(a + b) of all orderings; as a count,
-- (a + b) · |ballotSequences(a, b)| = (a − b) · (a + b choose a).
-- 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.Ballot
import List
import BinomialMore

-- A ballot is 0 (a vote for A) or 1 (a vote for B); an ordering of the
-- ballots is a list over 2, read in the order the votes are counted.
def zeroVotes (xs : Set) : Set := listCount{2}(0, xs)
def oneVotes (xs : Set) : Set := listCount{2}(1, xs)

theorem zero.in_two
  shows 0 ∈ 2
proof
  typecheck

theorem one.in_two
  shows 1 ∈ 2
proof
  typecheck

theorem zeroVotes.nil
  shows zeroVotes([]) = 0
proof
  unfold zeroVotes
  exact listCount.nil at 2, 0

theorem oneVotes.nil
  shows oneVotes([]) = 0
proof
  unfold oneVotes
  exact listCount.nil at 2, 1

theorem zeroVotes.natural (xs : List(2))
  shows Nat(zeroVotes(xs))
proof
  unfold zeroVotes
  exact listCount.natural at 2, 0, xs with *

theorem oneVotes.natural (xs : List(2))
  shows Nat(oneVotes(xs))
proof
  unfold oneVotes
  exact listCount.natural at 2, 1, xs with *

theorem zeroVotes.cons_zero (ys : List(2))
  shows zeroVotes(0 :: ys) = zeroVotes(ys) + 1
proof
  have zeroIn := zero.in_two
  have natural := listCount.natural at 2, 0, ys with *
  unfold zeroVotes
  rw listCount.cons_self
  ring

theorem oneVotes.cons_zero (ys : List(2))
  shows oneVotes(0 :: ys) = oneVotes(ys)
proof
  have zeroIn := zero.in_two
  have different := one.ne_zero
  unfold oneVotes
  exact listCount.cons_other with zeroIn, different

theorem zeroVotes.cons_one (ys : List(2))
  shows zeroVotes(1 :: ys) = zeroVotes(ys)
proof
  have oneIn := one.in_two
  have different := zero.ne_one
  unfold zeroVotes
  exact listCount.cons_other with oneIn, different

theorem oneVotes.cons_one (ys : List(2))
  shows oneVotes(1 :: ys) = oneVotes(ys) + 1
proof
  have oneIn := one.in_two
  have natural := listCount.natural at 2, 1, ys with *
  unfold oneVotes
  rw listCount.cons_self
  ring

-- Over the alphabet 2 every entry is a zero or a one.
theorem zeroVotes.add_ones (xs : List(2))
  shows zeroVotes(xs) + oneVotes(xs) = length{2}(xs)
proof
  induction xs
  | List.nil =>
    rw [zeroVotes.nil, oneVotes.nil, length.nil]
    compute
  | List.cons(x, tail) using ih =>
    have element : x ∈ 2 by assumption
    have zerosNat := zeroVotes.natural at tail with *
    have onesNat := oneVotes.natural at tail with *
    rw length.cons
    have alternatives := membership.two.forward with element
    cases alternatives with isZero | isOne
    . rw [isZero, zeroVotes.cons_zero, oneVotes.cons_zero, ← ih]
      ring
    . rw [isOne, zeroVotes.cons_one, oneVotes.cons_one, ← ih]
      ring

def ballots (n : Nat) → Set
  | zero => {List.nil}
  | succ(k) => consImage(0, ballots(k)) ∪ consImage(1, ballots(k))

theorem ballots.zero
  shows ballots(0) = {List.nil}
proof
  simp [ballots]

theorem ballots.succ (k : Nat)
  shows ballots(succ(k)) = consImage(0, ballots(k)) ∪ consImage(1, ballots(k))
proof
  simp [ballots]

theorem ballots.membership (n : Nat) (xs : Set)
  shows xs ∈ ballots(n) ↔ List(2, xs) ∧ length{2}(xs) = n
proof
  induction n generalizing xs
  | zero =>
    intro xs
    rw ballots.zero
    constructor
    . intro hx
      have equal := singleton.membership.forward with hx
      rw equal
      constructor
      . typecheck
      . exact length.nil at 2
    . intro hx
      obtain ⟨list, zero⟩ := hx
      have equal := length.eq_zero at 2, xs with list, zero
      rw equal
      apply singleton.membership.backward
      rfl
  | succ(k) using ih =>
    intro xs
    rw ballots.succ
    have zeroIn := zero.in_two
    have oneIn := one.in_two
    constructor
    . intro hx
      have split := union2.membership.forward with hx
      cases split with inZero | inOne
      . obtain ⟨ys, hy, shape⟩ := consImage.membership.forward with inZero
        have parts := ih with ys
        obtain ⟨list, len⟩ := parts.forward with hy
        rw ← shape
        constructor
        . typecheck
        . rw [length.cons, len]
          rfl
      . obtain ⟨ys, hy, shape⟩ := consImage.membership.forward with inOne
        have parts := ih with ys
        obtain ⟨list, len⟩ := parts.forward with hy
        rw ← shape
        constructor
        . typecheck
        . rw [length.cons, len]
          rfl
    . intro hx
      obtain ⟨list, len⟩ := hx
      have shape := List.cases at 2, xs with *
      cases shape with isNil | isCons
      . rw [isNil, length.nil] at len
        exfalso
        arith
      . obtain ⟨x, hx2, tail, tailList, xsShape⟩ := isCons
        rw [xsShape, length.cons] at len
        have tailLen := Naturals.injective.succ with len
        have parts := ih with tail
        have tailIn := parts.backward with ⟨tailList, tailLen⟩
        apply union2.membership.backward
        have alternatives := membership.two.forward with hx2
        cases alternatives with isZero | isOne
        . left
          apply consImage.membership.backward
          witness tail
          constructor
          . exact tailIn
          . rw [xsShape, isZero]
            rfl
        . right
          apply consImage.membership.backward
          witness tail
          constructor
          . exact tailIn
          . rw [xsShape, isOne]
            rfl

theorem ballots.finite (n : Nat)
  shows Finite(ballots(n))
proof
  induction n
  | zero =>
    rw ballots.zero
    exact Finite.singleton at List.nil
  | succ(k) using ih =>
    rw ballots.succ
    have left := consImage.finite at 0, ballots(k) with ih
    have right := consImage.finite at 1, ballots(k) with ih
    exact Finite.union with left, right

-- A stays ahead by more than k throughout: every nonempty prefix has fewer
-- ones than zeros plus k. AlwaysAhead is the case k = 0, the ballot
-- condition proper.
predicate AheadBy (k xs : Set) :=
  ∀ ys, List(2, ys) → ListPrefix(2, ys, xs) → ys ≠ [] →
    oneVotes(ys) < zeroVotes(ys) + k

predicate AlwaysAhead (xs : Set) :=
  ∀ ys, List(2, ys) → ListPrefix(2, ys, xs) → ys ≠ [] →
    oneVotes(ys) < zeroVotes(ys)

theorem AlwaysAhead.iff (xs : Set)
  shows AlwaysAhead(xs) ↔ AheadBy(0, xs)
proof
  constructor
  . intro ahead
    unfold AlwaysAhead at ahead
    unfold AheadBy
    intros ys list prefix nonempty
    have bound := ahead with ys, list, prefix, nonempty
    have zerosNat := zeroVotes.natural at ys with *
    have shape : zeroVotes(ys) + 0 = zeroVotes(ys) := add.right_unit at zeroVotes(ys)
    rw shape
    exact bound
  . intro ahead
    unfold AheadBy at ahead
    unfold AlwaysAhead
    intros ys list prefix nonempty
    have bound := ahead with ys, list, prefix, nonempty
    have zerosNat := zeroVotes.natural at ys with *
    have shape : zeroVotes(ys) + 0 = zeroVotes(ys) := add.right_unit at zeroVotes(ys)
    rw shape at bound
    exact bound

theorem AheadBy.nil (k : Nat)
  shows AheadBy(k, [])
proof
  unfold AheadBy
  intros ys list prefix nonempty
  obtain ⟨zs, zsList, equal⟩ := prefix
  symm at equal
  have parts := append.eq_nil with equal
  exfalso
  exact nonempty with parts.left

theorem AheadBy.cons_zero (k : Nat) (ys : List(2))
  shows AheadBy(k, 0 :: ys) ↔ AheadBy(k + 1, ys)
proof
  have zeroIn := zero.in_two
  constructor
  . intro ahead
    unfold AheadBy at ahead
    unfold AheadBy
    intros zs list prefix nonempty
    have prefixCons := ListPrefix.cons with zeroIn, prefix
    have nonemptyCons := List.cons_ne_nil at 2, 0, zs with *
    have consList : List(2, 0 :: zs) by typecheck
    have bound := ahead with 0 :: zs, consList, prefixCons, nonemptyCons
    rw [oneVotes.cons_zero, zeroVotes.cons_zero] at bound
    have zerosNat := zeroVotes.natural at zs with *
    have shape : zeroVotes(zs) + 1 + k = zeroVotes(zs) + (k + 1) by ring
    rw shape at bound
    exact bound
  . intro ahead
    unfold AheadBy at ahead
    unfold AheadBy
    intros zs list prefix nonempty
    have shape := List.cases at 2, zs with *
    cases shape with isNil | isCons
    . exfalso
      exact nonempty with isNil
    . obtain ⟨z, hz, tail, tailList, zsShape⟩ := isCons
      rw zsShape at prefix
      obtain ⟨rest, restList, consShape, tailPrefix⟩ :=
        ListPrefix.of_cons with hz, prefix
      obtain ⟨zIs, ysIs⟩ := List.inductive.injective.cons with consShape
      rw ← ysIs at tailPrefix
      rw [zsShape, ← zIs, oneVotes.cons_zero, zeroVotes.cons_zero]
      have zerosNat := zeroVotes.natural at tail with *
      have onesNat := oneVotes.natural at tail with *
      have tailShape := List.cases at 2, tail with *
      cases tailShape with tailNil | tailCons
      . rw [tailNil, oneVotes.nil, zeroVotes.nil]
        arith
      . obtain ⟨w, hw, more, moreList, tailIs⟩ := tailCons
        have tailNonempty : tail ≠ []
          rw tailIs
          exact List.cons_ne_nil at 2, w, more with *
        have bound := ahead with tail, tailList, tailPrefix, tailNonempty
        have reshape : zeroVotes(tail) + (k + 1) = zeroVotes(tail) + 1 + k by ring
        rw reshape at bound
        exact bound

theorem AheadBy.cons_one (k : Nat) (ys : List(2))
  shows AheadBy(k + 2, 1 :: ys) ↔ AheadBy(k + 1, ys)
proof
  have oneIn := one.in_two
  constructor
  . intro ahead
    unfold AheadBy at ahead
    unfold AheadBy
    intros zs list prefix nonempty
    have prefixCons := ListPrefix.cons with oneIn, prefix
    have nonemptyCons := List.cons_ne_nil at 2, 1, zs with *
    have consList : List(2, 1 :: zs) by typecheck
    have bound := ahead with 1 :: zs, consList, prefixCons, nonemptyCons
    rw [oneVotes.cons_one, zeroVotes.cons_one] at bound
    have zerosNat := zeroVotes.natural at zs with *
    have onesNat := oneVotes.natural at zs with *
    arith
  . intro ahead
    unfold AheadBy at ahead
    unfold AheadBy
    intros zs list prefix nonempty
    have shape := List.cases at 2, zs with *
    cases shape with isNil | isCons
    . exfalso
      exact nonempty with isNil
    . obtain ⟨z, hz, tail, tailList, zsShape⟩ := isCons
      rw zsShape at prefix
      obtain ⟨rest, restList, consShape, tailPrefix⟩ :=
        ListPrefix.of_cons with hz, prefix
      obtain ⟨zIs, ysIs⟩ := List.inductive.injective.cons with consShape
      rw ← ysIs at tailPrefix
      rw [zsShape, ← zIs, oneVotes.cons_one, zeroVotes.cons_one]
      have zerosNat := zeroVotes.natural at tail with *
      have onesNat := oneVotes.natural at tail with *
      have tailShape := List.cases at 2, tail with *
      cases tailShape with tailNil | tailCons
      . rw [tailNil, oneVotes.nil, zeroVotes.nil]
        arith
      . obtain ⟨w, hw, more, moreList, tailIs⟩ := tailCons
        have tailNonempty : tail ≠ []
          rw tailIs
          exact List.cons_ne_nil at 2, w, more with *
        have bound := ahead with tail, tailList, tailPrefix, tailNonempty
        arith

theorem AheadBy.not_cons_one (k : Nat) (ys : List(2))
  requires low : k ≤ 1
  shows ¬AheadBy(k, 1 :: ys)
proof
  have oneIn := one.in_two
  intro ahead
  unfold AheadBy at ahead
  have prefixNil := ListPrefix.nil at 2, ys with *
  have prefixOne := ListPrefix.cons with oneIn, prefixNil
  have nonempty := List.cons_ne_nil at 2, 1, [] with *
  have oneList : List(2, 1 :: []) by typecheck
  have bound := ahead with 1 :: [], oneList, prefixOne, nonempty
  have nilList : List(2, []) by typecheck
  have onesOne := oneVotes.cons_one at [] with nilList
  have zerosOne := zeroVotes.cons_one at [] with nilList
  rw [onesOne, zerosOne, oneVotes.nil, zeroVotes.nil] at bound
  arith

-- The orderings with a zeros and b ones along which A stays ahead by more
-- than k, and the ballot sequences proper.
def leadSequences (k a b : Nat) : Set :=
  { xs ∈ ballots(a + b) | zeroVotes(xs) = a ∧ oneVotes(xs) = b ∧ AheadBy(k, xs) }

def ballotSequences (a b : Nat) : Set :=
  { xs ∈ ballots(a + b) | zeroVotes(xs) = a ∧ AlwaysAhead(xs) }

theorem leadSequences.finite (k a b : Nat)
  shows Finite(leadSequences(k, a, b))
proof
  have subset : leadSequences(k, a, b) ⊆ ballots(a + b)
    intros xs hx
    have parts := leadSequences.membership.forward with hx
    exact parts.left
  exact Finite.subset with (ballots.finite at a + b), subset

theorem ballotSequences.finite (a b : Nat)
  shows Finite(ballotSequences(a, b))
proof
  have subset : ballotSequences(a, b) ⊆ ballots(a + b)
    intros xs hx
    have parts := ballotSequences.membership.forward with hx
    exact parts.left
  exact Finite.subset with (ballots.finite at a + b), subset

-- The orderings in leadSequences(k, a, b) whose first ballot is 0, and
-- those whose first ballot is 1.
def headZero (k a b : Nat) : Set :=
  { xs ∈ leadSequences(k, a, b) | ∃ ys, 0 :: ys = xs }

def headOne (k a b : Nat) : Set :=
  { xs ∈ leadSequences(k, a, b) | ∃ ys, 1 :: ys = xs }

theorem headZero.subset (k a b : Nat)
  shows headZero(k, a, b) ⊆ leadSequences(k, a, b)
proof
  intros xs hx
  have parts := headZero.membership.forward with hx
  exact parts.left

theorem headOne.subset (k a b : Nat)
  shows headOne(k, a, b) ⊆ leadSequences(k, a, b)
proof
  intros xs hx
  have parts := headOne.membership.forward with hx
  exact parts.left

-- A nonempty ordering begins with 0 or with 1, so the count splits.
theorem leadSequences.split (k a b : Nat)
  requires positive : 0 < a + b
  shows cardinality(leadSequences(k, a, b)) =
    cardinality(headZero(k, a, b)) + cardinality(headOne(k, a, b))
proof
  have finite := leadSequences.finite at k, a, b
  have zeroSubset := headZero.subset at k, a, b
  have oneSubset := headOne.subset at k, a, b
  have finiteZero := Finite.subset with finite, zeroSubset
  have finiteOne := Finite.subset with finite, oneSubset
  have disjoint : Disjoint(headZero(k, a, b), headOne(k, a, b))
    apply Disjoint.iff.backward
    intros xs hx inOne
    have zeroParts := headZero.membership.forward with hx
    have oneParts := headOne.membership.forward with inOne
    obtain ⟨ys, ysShape⟩ := zeroParts.right
    obtain ⟨zs, zsShape⟩ := oneParts.right
    rw ← ysShape at zsShape
    obtain ⟨headEqual, tailEqual⟩ :=
      List.inductive.injective.cons with zsShape
    exact one.ne_zero with headEqual
  have union : leadSequences(k, a, b) = headZero(k, a, b) ∪ headOne(k, a, b)
    ext xs
    constructor
    . intro hx
      have parts := leadSequences.membership.forward with hx
      obtain ⟨list, len⟩ := ballots.membership.forward with parts.left
      have shape := List.cases at 2, xs with *
      apply union2.membership.backward
      cases shape with isNil | isCons
      . exfalso
        rw [isNil, length.nil] at len
        arith
      . obtain ⟨x, hx2, tail, tailList, xsShape⟩ := isCons
        have alternatives := membership.two.forward with hx2
        cases alternatives with isZero | isOne
        . left
          apply headZero.membership.backward
          constructor
          . exact hx
          . witness tail
            rw [xsShape, isZero]
            rfl
        . right
          apply headOne.membership.backward
          constructor
          . exact hx
          . witness tail
            rw [xsShape, isOne]
            rfl
    . intro hx
      have split := union2.membership.forward with hx
      unfold Subset at zeroSubset
      unfold Subset at oneSubset
      cases split with inZero | inOne
      . exact zeroSubset with xs, inZero
      . exact oneSubset with xs, inOne
  rw union
  exact cardinality.union_disjoint with finiteZero, finiteOne, disjoint

-- Removing a first ballot 0 lowers the number of zeros by one and raises
-- the lead by one.
theorem headZero.succ (k a b : Nat)
  shows headZero(k, a + 1, b) = consImage(0, leadSequences(k + 1, a, b))
proof
  have zeroIn := zero.in_two
  ext xs
  constructor
  . intro hx
    have parts := headZero.membership.forward with hx
    obtain ⟨ys, ysShape⟩ := parts.right
    obtain ⟨inBallots, zerosIs, onesIs, ahead⟩ :=
      leadSequences.membership.forward with parts.left
    obtain ⟨list, len⟩ := ballots.membership.forward with inBallots
    rw ← ysShape at list
    obtain ⟨headIn, ysList⟩ := List.cons_parts with list
    rw ← ysShape at len
    rw ← ysShape at zerosIs
    rw ← ysShape at onesIs
    rw ← ysShape at ahead
    rw [length.cons] at len
    rw zeroVotes.cons_zero at zerosIs
    rw oneVotes.cons_zero at onesIs
    have iff := AheadBy.cons_zero at k, ys with *
    have zerosNat := zeroVotes.natural at ys with *
    apply consImage.membership.backward
    witness ys
    constructor
    . apply leadSequences.membership.backward
      refine ⟨?_, ?_, ?_, ?_⟩
      . apply ballots.membership.backward
        constructor
        . exact ysList
        . have shape : a + 1 + b = succ(a + b) by ring
          rw shape at len
          exact Naturals.injective.succ with len
      . exact add.right_cancel at zeroVotes(ys), a, 1 with zerosIs
      . exact onesIs
      . exact iff.forward with ahead
    . exact ysShape
  . intro hx
    obtain ⟨ys, hy, ysShape⟩ := consImage.membership.forward with hx
    obtain ⟨inBallots, zerosIs, onesIs, ahead⟩ :=
      leadSequences.membership.forward with hy
    obtain ⟨ysList, len⟩ := ballots.membership.forward with inBallots
    have iff := AheadBy.cons_zero at k, ys with *
    apply headZero.membership.backward
    constructor
    . apply leadSequences.membership.backward
      rw ← ysShape
      refine ⟨?_, ?_, ?_, ?_⟩
      . apply ballots.membership.backward
        constructor
        . typecheck
        . rw [length.cons, len]
          ring
      . rw [zeroVotes.cons_zero, zerosIs]
        rfl
      . rw oneVotes.cons_zero
        exact onesIs
      . exact iff.backward with ahead
    . witness ys
      exact ysShape

theorem headZero.zero (k b : Nat)
  shows headZero(k, 0, b) = ∅
proof
  ext xs
  constructor
  . intro hx
    exfalso
    have parts := headZero.membership.forward with hx
    obtain ⟨ys, ysShape⟩ := parts.right
    obtain ⟨inBallots, zerosIs, onesIs, ahead⟩ :=
      leadSequences.membership.forward with parts.left
    obtain ⟨list, len⟩ := ballots.membership.forward with inBallots
    rw ← ysShape at list
    obtain ⟨headIn, ysList⟩ := List.cons_parts with list
    rw [← ysShape, zeroVotes.cons_zero] at zerosIs
    have zerosNat := zeroVotes.natural at ys with *
    arith
  . intro hx
    exact empty.subset at headZero(k, 0, b) with xs, hx

-- Removing a first ballot 1 lowers the number of ones and the lead by one.
theorem headOne.succ (k a b : Nat)
  shows headOne(k + 2, a, b + 1) = consImage(1, leadSequences(k + 1, a, b))
proof
  have oneIn := one.in_two
  ext xs
  constructor
  . intro hx
    have parts := headOne.membership.forward with hx
    obtain ⟨ys, ysShape⟩ := parts.right
    obtain ⟨inBallots, zerosIs, onesIs, ahead⟩ :=
      leadSequences.membership.forward with parts.left
    obtain ⟨list, len⟩ := ballots.membership.forward with inBallots
    rw ← ysShape at list
    obtain ⟨headIn, ysList⟩ := List.cons_parts with list
    rw ← ysShape at len
    rw ← ysShape at zerosIs
    rw ← ysShape at onesIs
    rw ← ysShape at ahead
    rw [length.cons] at len
    rw zeroVotes.cons_one at zerosIs
    rw oneVotes.cons_one at onesIs
    have iff := AheadBy.cons_one at k, ys with *
    have onesNat := oneVotes.natural at ys with *
    apply consImage.membership.backward
    witness ys
    constructor
    . apply leadSequences.membership.backward
      refine ⟨?_, ?_, ?_, ?_⟩
      . apply ballots.membership.backward
        constructor
        . exact ysList
        . have shape : a + (b + 1) = succ(a + b) by ring
          rw shape at len
          exact Naturals.injective.succ with len
      . exact zerosIs
      . exact add.right_cancel at oneVotes(ys), b, 1 with onesIs
      . exact iff.forward with ahead
    . exact ysShape
  . intro hx
    obtain ⟨ys, hy, ysShape⟩ := consImage.membership.forward with hx
    obtain ⟨inBallots, zerosIs, onesIs, ahead⟩ :=
      leadSequences.membership.forward with hy
    obtain ⟨ysList, len⟩ := ballots.membership.forward with inBallots
    have iff := AheadBy.cons_one at k, ys with *
    apply headOne.membership.backward
    constructor
    . apply leadSequences.membership.backward
      rw ← ysShape
      refine ⟨?_, ?_, ?_, ?_⟩
      . apply ballots.membership.backward
        constructor
        . typecheck
        . rw [length.cons, len]
          ring
      . rw zeroVotes.cons_one
        exact zerosIs
      . rw [oneVotes.cons_one, onesIs]
        rfl
      . exact iff.backward with ahead
    . witness ys
      exact ysShape

-- With a lead of at most one, no ordering may begin with a ballot 1.
theorem headOne.low (k a b : Nat)
  requires low : k ≤ 1
  shows headOne(k, a, b) = ∅
proof
  ext xs
  constructor
  . intro hx
    exfalso
    have parts := headOne.membership.forward with hx
    obtain ⟨ys, ysShape⟩ := parts.right
    obtain ⟨inBallots, zerosIs, onesIs, ahead⟩ :=
      leadSequences.membership.forward with parts.left
    obtain ⟨list, len⟩ := ballots.membership.forward with inBallots
    rw ← ysShape at list
    obtain ⟨headIn, ysList⟩ := List.cons_parts with list
    rw ← ysShape at ahead
    have blocked := AheadBy.not_cons_one at k, ys with *
    exact blocked with ahead
  . intro hx
    exact empty.subset at headOne(k, a, b) with xs, hx

theorem headOne.zero (k a : Nat)
  shows headOne(k, a, 0) = ∅
proof
  ext xs
  constructor
  . intro hx
    exfalso
    have parts := headOne.membership.forward with hx
    obtain ⟨ys, ysShape⟩ := parts.right
    obtain ⟨inBallots, zerosIs, onesIs, ahead⟩ :=
      leadSequences.membership.forward with parts.left
    obtain ⟨list, len⟩ := ballots.membership.forward with inBallots
    rw ← ysShape at list
    obtain ⟨headIn, ysList⟩ := List.cons_parts with list
    rw [← ysShape, oneVotes.cons_one] at onesIs
    have onesNat := oneVotes.natural at ys with *
    arith
  . intro hx
    exact empty.subset at headOne(k, a, 0) with xs, hx

-- The empty ordering is the only one without ballots.
theorem leadSequences.nil (k : Nat)
  shows leadSequences(k, 0, 0) = {[]}
proof
  have nilList : List(2, []) by typecheck
  have shape : (0:N) + 0 = 0 by compute
  ext xs
  constructor
  . intro hx
    obtain ⟨inBallots, zerosIs, onesIs, ahead⟩ :=
      leadSequences.membership.forward with hx
    obtain ⟨list, len⟩ := ballots.membership.forward with inBallots
    rw shape at len
    have equal := length.eq_zero at 2, xs with list, len
    rw equal
    apply singleton.membership.backward
    rfl
  . intro hx
    have equal := singleton.membership.forward with hx
    rw equal
    apply leadSequences.membership.backward
    refine ⟨?_, zeroVotes.nil, oneVotes.nil, AheadBy.nil at k⟩
    apply ballots.membership.backward
    rw [shape, length.nil]
    exact ⟨nilList, rfl⟩

theorem leadSequences.nil_count (k : Nat)
  shows cardinality(leadSequences(k, 0, 0)) = 1
proof
  rw leadSequences.nil
  exact cardinality.singleton at []

-- The reflection count: among the orderings with a zeros and b ones, those
-- along which A stays ahead by more than k number
-- (n choose a) − (n choose a + k + 1), stated without subtraction.
theorem leadSequences.closed_form (n k a b : Nat)
  requires total : a + b = n
  requires bound : b ≤ a + k
  shows cardinality(leadSequences(k + 1, a, b)) + (n choose (a + k + 1)) =
    n choose a
proof
  induction n generalizing k a b
  | zero =>
    intros k a b kNat aNat bNat total bound
    obtain ⟨aZero, bZero⟩ := add.eq_zero at b, a with total
    rw [aZero, bZero, leadSequences.nil_count]
    have shape : (0:N) + k + 1 = succ(k) by ring
    rw [shape, choose.zero_left, choose.zero_right]
    compute
  | succ(m) using ih =>
    intros k a b kNat aNat bNat total bound
    have positive : 0 < a + b by arith
    have split := leadSequences.split at k + 1, a, b with positive
    have aShape := natural.zero_or_successor with aNat
    have bShape := natural.zero_or_successor with bNat
    have kShape := natural.zero_or_successor with kNat
    cases aShape with aZero | aSucc
    . -- No zeroVotes: only the all-oneVotes ordering, which needs a lead above b.
      rw aZero at total
      rw aZero at bound
      rw aZero at split
      rw aZero
      cases bShape with bZero | bSucc
      . exfalso
        rw bZero at total
        arith
      . obtain ⟨b2, b2Nat, rfl⟩ := bSucc
        cases kShape with kZero | kSucc
        . exfalso
          rw kZero at bound
          arith
        . obtain ⟨k2, k2Nat, rfl⟩ := kSucc
          have mIs : b2 = m
            have shape : (0:N) + succ(b2) = succ(b2) :=
              add.left_unit at succ(b2)
            rw shape at total
            exact Naturals.injective.succ with total
          have zeroCount : cardinality(headZero(succ(k2) + 1, 0, succ(b2))) = 0
            rw headZero.zero
            exact cardinality.empty
          have oneShape : headOne(succ(k2) + 1, 0, succ(b2)) =
              consImage(1, leadSequences(k2 + 1, 0, b2))
            have kForm : succ(k2) + 1 = k2 + 2 by ring
            have bForm : succ(b2) = b2 + 1 by ring
            rw [kForm, bForm]
            exact headOne.succ at k2, 0, b2
          have oneCount : cardinality(headOne(succ(k2) + 1, 0, succ(b2))) =
              cardinality(leadSequences(k2 + 1, 0, b2))
            rw oneShape
            exact consImage.cardinality with
              (leadSequences.finite at k2 + 1, 0, b2)
          have zeroNat : Nat(0) by typecheck
          have total2 : (0:N) + b2 = m
            rw ← mIs
            exact add.left_unit at b2
          have bound2 : b2 ≤ 0 + k2 by arith
          have inner := ih with k2, 0, b2, k2Nat, zeroNat, b2Nat, total2, bound2
          have vanish : m choose (0 + k2 + 1) = 0
            have form : (0:N) + k2 + 1 = succ(k2) by ring
            have le : m ≤ k2 by arith
            rw form
            exact choose.vanishes with le
          have vanishNext : succ(m) choose (0 + succ(k2) + 1) = 0
            have form : (0:N) + succ(k2) + 1 = succ(succ(k2)) by ring
            have le : succ(m) ≤ succ(k2) by arith
            rw form
            exact choose.vanishes with le
          have cardNat := cardinality.nat with
            (leadSequences.finite at k2 + 1, 0, b2)
          rw [vanish, choose.zero_right] at inner
          rw [split, zeroCount, oneCount, vanishNext, choose.zero_right]
          ring [inner]
    . obtain ⟨a2, a2Nat, rfl⟩ := aSucc
      cases bShape with bZero | bSucc
      . -- No oneVotes: only the all-zeroVotes ordering.
        rw bZero at total
        rw bZero at bound
        rw bZero at split
        rw bZero
        have saNat : Nat(succ(a2)) by typecheck
        have mIs : a2 = m
          have shape : succ(a2) + 0 = succ(a2) := add.right_unit at succ(a2)
          rw shape at total
          exact Naturals.injective.succ with total
        have oneCount : cardinality(headOne(k + 1, succ(a2), 0)) = 0
          rw headOne.zero
          exact cardinality.empty
        have zeroShape : headZero(k + 1, succ(a2), 0) =
            consImage(0, leadSequences(k + 1 + 1, a2, 0))
          have aForm : succ(a2) = a2 + 1 by ring
          rw aForm
          exact headZero.succ at k + 1, a2, 0
        have zeroCount : cardinality(headZero(k + 1, succ(a2), 0)) =
            cardinality(leadSequences(k + 1 + 1, a2, 0))
          rw zeroShape
          exact consImage.cardinality with
            (leadSequences.finite at k + 1 + 1, a2, 0)
        have zeroNat : Nat(0) by typecheck
        have kNat2 : Nat(k + 1) by typecheck
        have total2 : a2 + 0 = m
          rw ← mIs
          exact add.right_unit at a2
        have bound2 : 0 ≤ a2 + (k + 1) by arith
        have inner := ih with k + 1, a2, 0, kNat2, a2Nat, zeroNat, total2,
          bound2
        have vanish : m choose (a2 + (k + 1) + 1) = 0
          have form : a2 + (k + 1) + 1 = succ(a2 + (k + 1)) by ring
          have le : m ≤ a2 + (k + 1) by arith
          rw form
          exact choose.vanishes with le
        have vanishNext : succ(m) choose (succ(a2) + k + 1) = 0
          have form : succ(a2) + k + 1 = succ(succ(a2) + k) by ring
          have le : succ(m) ≤ succ(a2) + k by arith
          rw form
          exact choose.vanishes with le
        have diagonal : m choose a2 = 1
          rw mIs
          exact choose.diagonal at m
        have diagonalNext : succ(m) choose succ(a2) = 1
          rw mIs
          exact choose.diagonal at succ(m)
        have cardNat := cardinality.nat with
          (leadSequences.finite at k + 1 + 1, a2, 0)
        rw [vanish, diagonal] at inner
        rw [split, zeroCount, oneCount, vanishNext, diagonalNext]
        ring [inner]
      . obtain ⟨b2, b2Nat, rfl⟩ := bSucc
        cases kShape with kZero | kSucc
        . -- Lead one: the first ballot must be 0.
          rw kZero at bound
          rw kZero at split
          rw kZero
          have mIs : a2 + succ(b2) = m
            have shape : succ(a2) + succ(b2) = succ(a2 + succ(b2)) by ring
            rw shape at total
            exact Naturals.injective.succ with total
          have oneCount : cardinality(headOne(0 + 1, succ(a2), succ(b2))) = 0
            have low : (0:N) + 1 ≤ 1 by arith
            have empty := headOne.low at 0 + 1, succ(a2), succ(b2) with low
            rw empty
            exact cardinality.empty
          have zeroShape : headZero(0 + 1, succ(a2), succ(b2)) =
              consImage(0, leadSequences(0 + 1 + 1, a2, succ(b2)))
            have aForm : succ(a2) = a2 + 1 by ring
            rw aForm
            exact headZero.succ at 0 + 1, a2, succ(b2)
          have zeroCount : cardinality(headZero(0 + 1, succ(a2), succ(b2))) =
              cardinality(leadSequences(0 + 1 + 1, a2, succ(b2)))
            rw zeroShape
            exact consImage.cardinality with
              (leadSequences.finite at 0 + 1 + 1, a2, succ(b2))
          have oneNat : Nat(0 + 1) by typecheck
          have sbNat : Nat(succ(b2)) by typecheck
          have bound2 : succ(b2) ≤ a2 + (0 + 1) by arith
          have inner := ih with 0 + 1, a2, succ(b2), oneNat, a2Nat, sbNat, mIs,
            bound2
          have pascalA : succ(m) choose succ(a2) =
              (m choose a2) + (m choose succ(a2)) := choose.succ_succ at m, a2
          have pascalB : succ(m) choose (succ(a2) + 0 + 1) =
              (m choose succ(a2)) + (m choose succ(succ(a2)))
            have form : succ(a2) + 0 + 1 = succ(succ(a2)) by ring
            rw form
            exact choose.succ_succ at m, succ(a2)
          have innerForm : a2 + (0 + 1) + 1 = succ(succ(a2)) by ring
          rw innerForm at inner
          have cardNat := cardinality.nat with
            (leadSequences.finite at 0 + 1 + 1, a2, succ(b2))
          have firstNat := choose.natural at m, a2
          have secondNat := choose.natural at m, succ(a2)
          have thirdNat := choose.natural at m, succ(succ(a2))
          rw [split, zeroCount, oneCount, pascalB, pascalA]
          ring [inner]
        . -- Both first ballots are possible.
          obtain ⟨k2, k2Nat, rfl⟩ := kSucc
          have mIs : a2 + succ(b2) = m
            have shape : succ(a2) + succ(b2) = succ(a2 + succ(b2)) by ring
            rw shape at total
            exact Naturals.injective.succ with total
          have mIs2 : succ(a2) + b2 = m
            have shape : succ(a2) + succ(b2) = succ(succ(a2) + b2) by ring
            rw shape at total
            exact Naturals.injective.succ with total
          have zeroShape : headZero(succ(k2) + 1, succ(a2), succ(b2)) =
              consImage(0, leadSequences(succ(k2) + 1 + 1, a2, succ(b2)))
            have aForm : succ(a2) = a2 + 1 by ring
            rw aForm
            exact headZero.succ at succ(k2) + 1, a2, succ(b2)
          have zeroCount :
              cardinality(headZero(succ(k2) + 1, succ(a2), succ(b2))) =
              cardinality(leadSequences(succ(k2) + 1 + 1, a2, succ(b2)))
            rw zeroShape
            exact consImage.cardinality with
              (leadSequences.finite at succ(k2) + 1 + 1, a2, succ(b2))
          have oneShape : headOne(succ(k2) + 1, succ(a2), succ(b2)) =
              consImage(1, leadSequences(k2 + 1, succ(a2), b2))
            have kForm : succ(k2) + 1 = k2 + 2 by ring
            have bForm : succ(b2) = b2 + 1 by ring
            rw [kForm, bForm]
            exact headOne.succ at k2, succ(a2), b2
          have oneCount :
              cardinality(headOne(succ(k2) + 1, succ(a2), succ(b2))) =
              cardinality(leadSequences(k2 + 1, succ(a2), b2))
            rw oneShape
            exact consImage.cardinality with
              (leadSequences.finite at k2 + 1, succ(a2), b2)
          have kNat2 : Nat(succ(k2) + 1) by typecheck
          have sbNat : Nat(succ(b2)) by typecheck
          have saNat : Nat(succ(a2)) by typecheck
          have bound0 : succ(b2) ≤ a2 + (succ(k2) + 1) by arith
          have bound1 : b2 ≤ succ(a2) + k2 by arith
          have inner0 := ih with succ(k2) + 1, a2, succ(b2), kNat2, a2Nat,
            sbNat, mIs, bound0
          have inner1 := ih with k2, succ(a2), b2, k2Nat, saNat, b2Nat, mIs2,
            bound1
          have pascalA : succ(m) choose succ(a2) =
              (m choose a2) + (m choose succ(a2)) := choose.succ_succ at m, a2
          have pascalB : succ(m) choose (succ(a2) + succ(k2) + 1) =
              (m choose (succ(a2) + k2 + 1)) +
              (m choose succ(succ(a2) + k2 + 1))
            have form : succ(a2) + succ(k2) + 1 = succ(succ(a2) + k2 + 1)
              ring
            rw form
            exact choose.succ_succ at m, succ(a2) + k2 + 1
          have innerForm : a2 + (succ(k2) + 1) + 1 = succ(succ(a2) + k2 + 1)
            ring
          rw innerForm at inner0
          have cardNat0 := cardinality.nat with
            (leadSequences.finite at succ(k2) + 1 + 1, a2, succ(b2))
          have cardNat1 := cardinality.nat with
            (leadSequences.finite at k2 + 1, succ(a2), b2)
          have firstNat := choose.natural at m, a2
          have secondNat := choose.natural at m, succ(a2)
          have thirdNat := choose.natural at m, succ(a2) + k2 + 1
          have fourthNat := choose.natural at m, succ(succ(a2) + k2 + 1)
          rw [split, zeroCount, oneCount, pascalB, pascalA]
          ring [inner0, inner1]

-- A ballot sequence begins with a vote for A, and what follows keeps A
-- ahead by more than one.
theorem ballotSequences.eq (a b : Nat)
  shows ballotSequences(a + 1, b) = consImage(0, leadSequences(1, a, b))
proof
  have zeroIn := zero.in_two
  have form : (0:N) + 1 = 1 by compute
  ext xs
  constructor
  . intro hx
    obtain ⟨inBallots, zerosIs, ahead⟩ :=
      ballotSequences.membership.forward with hx
    obtain ⟨list, len⟩ := ballots.membership.forward with inBallots
    have shape := List.cases at 2, xs with *
    cases shape with isNil | isCons
    . exfalso
      rw [isNil, length.nil] at len
      arith
    . obtain ⟨x, hx2, tail, tailList, xsShape⟩ := isCons
      have iff := AlwaysAhead.iff at xs
      have aheadBy := iff.forward with ahead
      rw xsShape at aheadBy
      have alternatives := membership.two.forward with hx2
      cases alternatives with isZero | isOne
      . rw isZero at xsShape
        rw isZero at aheadBy
        rw xsShape at len
        rw xsShape at zerosIs
        rw [length.cons] at len
        rw zeroVotes.cons_zero at zerosIs
        have zerosNat := zeroVotes.natural at tail with *
        have onesNat := oneVotes.natural at tail with *
        have tailLen : length{2}(tail) = a + b
          have reshape : a + 1 + b = succ(a + b) by ring
          rw reshape at len
          exact Naturals.injective.succ with len
        have tailZeros : zeroVotes(tail) = a :=
          add.right_cancel at zeroVotes(tail), a, 1 with zerosIs
        have tailOnes : oneVotes(tail) = b
          have sum := zeroVotes.add_ones at tail with *
          rw [tailLen, tailZeros] at sum
          exact add.left_cancel at oneVotes(tail), b, a with sum
        have consIff := AheadBy.cons_zero at 0, tail with *
        have tailAhead := consIff.forward with aheadBy
        rw form at tailAhead
        apply consImage.membership.backward
        witness tail
        constructor
        . apply leadSequences.membership.backward
          refine ⟨?_, tailZeros, tailOnes, tailAhead⟩
          apply ballots.membership.backward
          exact ⟨tailList, tailLen⟩
        . symm
          exact xsShape
      . exfalso
        rw isOne at aheadBy
        have low : (0:N) ≤ 1 by arith
        have blocked := AheadBy.not_cons_one at 0, tail with *
        exact blocked with aheadBy
  . intro hx
    obtain ⟨ys, hy, ysShape⟩ := consImage.membership.forward with hx
    obtain ⟨inBallots, zerosIs, onesIs, ahead⟩ :=
      leadSequences.membership.forward with hy
    obtain ⟨ysList, len⟩ := ballots.membership.forward with inBallots
    rw ← form at ahead
    have consIff := AheadBy.cons_zero at 0, ys with *
    have consAhead := consIff.backward with ahead
    have iff := AlwaysAhead.iff at 0 :: ys
    apply ballotSequences.membership.backward
    rw ← ysShape
    refine ⟨?_, ?_, iff.backward with consAhead⟩
    . apply ballots.membership.backward
      constructor
      . typecheck
      . rw [length.cons, len]
        ring
    . rw [zeroVotes.cons_zero, zerosIs]
      rfl

-- The count in additive form: with b ≤ a, the ballot sequences for a + 1
-- votes for A number (a + b choose a) − (a + b choose a + 1).
theorem ballotSequences.count_succ (a b : Nat)
  requires le : b ≤ a
  shows cardinality(ballotSequences(a + 1, b)) + ((a + b) choose (a + 1)) =
    (a + b) choose a
proof
  have zeroNat : Nat(0) by typecheck
  have total : a + b = a + b by rfl
  have bound : b ≤ a + 0 by arith
  have closed := leadSequences.closed_form at a + b, 0, a, b with total, bound
  have formK : (0:N) + 1 = 1 by compute
  have formA : a + 0 + 1 = a + 1 by ring
  rw [formK, formA] at closed
  rw [ballotSequences.eq, consImage.cardinality]
  exact closed

-- All orderings of the ballots: a zeros among a + b entries, counted by
-- the binomial coefficient through the lead family at a lead above b.
def orderings (a b : Nat) : Set := { xs ∈ ballots(a + b) | zeroVotes(xs) = a }

theorem orderings.finite (a b : Nat)
  shows Finite(orderings(a, b))
proof
  have subset : orderings(a, b) ⊆ ballots(a + b)
    intros xs hx
    have parts := orderings.membership.forward with hx
    exact parts.left
  exact Finite.subset with (ballots.finite at a + b), subset

theorem ballotSequences.subset_orderings (a b : Nat)
  shows ballotSequences(a, b) ⊆ orderings(a, b)
proof
  intros xs hx
  obtain ⟨inBallots, zerosIs, ahead⟩ :=
    ballotSequences.membership.forward with hx
  apply orderings.membership.backward
  exact ⟨inBallots, zerosIs⟩

theorem oneVotes.prefix_le (ys xs : List(2))
  requires prefix : ListPrefix(2, ys, xs)
  shows oneVotes(ys) ≤ oneVotes(xs)
proof
  have le := ListPrefix.count_le at 2, 1, ys, xs with prefix
  unfold oneVotes
  exact le

theorem orderings.eq (a b : Nat)
  shows orderings(a, b) = leadSequences(b + 1, a, b)
proof
  ext xs
  constructor
  . intro hx
    obtain ⟨inBallots, zerosIs⟩ := orderings.membership.forward with hx
    obtain ⟨list, len⟩ := ballots.membership.forward with inBallots
    have onesIs : oneVotes(xs) = b
      have sum := zeroVotes.add_ones at xs with *
      rw [len, zerosIs] at sum
      exact add.left_cancel at oneVotes(xs), b, a with sum
    apply leadSequences.membership.backward
    refine ⟨inBallots, zerosIs, onesIs, ?_⟩
    unfold AheadBy
    intros ys ysList prefix nonempty
    have le := oneVotes.prefix_le with prefix
    have zerosNat := zeroVotes.natural at ys with *
    have onesNat := oneVotes.natural at ys with *
    rw onesIs at le
    arith
  . intro hx
    obtain ⟨inBallots, zerosIs, onesIs, ahead⟩ :=
      leadSequences.membership.forward with hx
    apply orderings.membership.backward
    exact ⟨inBallots, zerosIs⟩

theorem orderings.count (a b : Nat)
  shows cardinality(orderings(a, b)) = (a + b) choose a
proof
  have total : a + b = a + b by rfl
  have bound : b ≤ a + b by arith
  have closed := leadSequences.closed_form at a + b, b, a, b with total, bound
  have vanish : (a + b) choose (a + b + 1) = 0
    have reflexive : a + b ≤ a + b by arith
    have shape : a + b + 1 = succ(a + b) by ring
    rw shape
    exact choose.vanishes with reflexive
  have cardNat := cardinality.nat with (leadSequences.finite at b + 1, a, b)
  rw [vanish, add.right_unit] at closed
  rw orderings.eq
  exact closed

-- Bertrand's ballot theorem, as a count.
theorem ballot.count (a b : Nat)
  requires lead : b < a
  shows (a + b) * cardinality(ballotSequences(a, b)) =
    (a - b) * ((a + b) choose a)
proof
  have aShape := natural.zero_or_successor at a with *
  cases aShape with aZero | aSucc
  . exfalso
    rw aZero at lead
    arith
  . obtain ⟨a2, a2Nat, rfl⟩ := aSucc
    have le : b ≤ a2 by arith
    have form : succ(a2) = a2 + 1 by ring
    rw form
    have count := ballotSequences.count_succ with le
    have absorb : (a2 + 1) * ((a2 + 1 + b) choose (a2 + 1)) =
        (a2 + 1 + b) * ((a2 + b) choose a2)
      have shape : a2 + 1 + b = a2 + b + 1 by ring
      rw shape
      exact choose.absorption_at at a2 + b, a2
    have finite := ballotSequences.finite at a2 + 1, b
    have cardNat := cardinality.nat with finite
    have xNat := choose.natural at a2 + b, a2 + 1
    have cNat := choose.natural at a2 + 1 + b, a2 + 1
    have main : (a2 + 1 + b) * cardinality(ballotSequences(a2 + 1, b)) +
        b * ((a2 + 1 + b) choose (a2 + 1)) =
        (a2 + 1) * ((a2 + 1 + b) choose (a2 + 1))
      have bShape := natural.zero_or_successor at b with *
      cases bShape with bZero | bSucc
      . rw bZero at count
        rw bZero at absorb
        rw bZero
        have vanish : (a2 + 0) choose (a2 + 1) = 0
          have reflexive : a2 ≤ a2 by arith
          have shape : a2 + 0 = a2 := add.right_unit at a2
          have shapeOne : a2 + 1 = succ(a2) := add.right_one at a2
          rw [shape, shapeOne]
          exact choose.vanishes with reflexive
        rw [absorb, ← count, vanish]
        ring
      . obtain ⟨b2, b2Nat, rfl⟩ := bSucc
        have complement : succ(b2) * ((a2 + 1 + succ(b2)) choose (a2 + 1)) =
            (a2 + 1 + succ(b2)) * ((a2 + succ(b2)) choose (a2 + 1))
          have le2 : a2 + 1 ≤ a2 + succ(b2) by arith
          have comp := choose.absorption_complement at a2 + succ(b2), a2 + 1
            with le2
          have diff : a2 + succ(b2) - (a2 + 1) = b2
            have shape : a2 + succ(b2) = a2 + 1 + b2 by ring
            rw shape
            exact sub.add_left at b2, a2 + 1
          have coefficient : b2 + 1 = succ(b2) by ring
          have shape : a2 + 1 + succ(b2) = a2 + succ(b2) + 1 by ring
          rw [diff, coefficient] at comp
          rw shape
          exact comp
        rw [absorb, ← count, complement]
        ring
    have le2 : b ≤ a2 + 1 by arith
    have split := mul.split_sub at a2 + 1, b, (a2 + 1 + b) choose (a2 + 1)
      with le2
    rw split at main
    exact add.right_cancel at
      (a2 + 1 + b) * cardinality(ballotSequences(a2 + 1, b)),
      (a2 + 1 - b) * ((a2 + 1 + b) choose (a2 + 1)),
      b * ((a2 + 1 + b) choose (a2 + 1)) with main