-- Binomial coefficients by Pascal's rule, and their factorial form
-- 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.NumberTheory.Binomial
import Nat

-- The row n = 0 of Pascal's triangle: 1 at k = 0, then 0.
def chooseZero (k : Nat) → Nat
  | zero => 1
  | succ(k) => 0

-- Pascal's rule at column k, given the two entries above it: the entry at
-- column 0 is 1, and the entry at column succ(j) is the sum of the entries at
-- columns j and succ(j) of the row above.
def pascalStep (k above left : Nat) → Nat
  | zero, above, left => 1
  | succ(j), above, left => add(left, above)

-- The number of k-element subsets of an n-element set, by Pascal's rule:
-- choose(n + 1, k + 1) = choose(n, k) + choose(n, k + 1).
def choose (n k : Nat) → Nat
  | zero, k => chooseZero(k)
  | succ(n), k => pascalStep(k, choose(n, k), choose(n, predecessor(k)))
  descends n

-- The binomial coefficient is written n choose k; it binds more loosely
-- than +, so 2 * n + 1 choose k is the coefficient of 2n + 1.
notation infixl 55 "choose" => choose for Nat

theorem choose.zero_right (n : Nat)
  shows choose(n, 0) = 1
proof
  induction n
  | zero => compute
  | succ(k) using ih => settle with choose, pascalStep

theorem choose.zero_left (k : Nat)
  shows choose(0, succ(k)) = 0
proof
  simp [choose, chooseZero]

theorem choose.succ_succ (n k : Nat)
  shows choose(succ(n), succ(k)) = choose(n, k) + choose(n, succ(k))
proof
  simp [choose, pascalStep, predecessor]

-- Nothing is chosen above the diagonal.
theorem choose.above (n : Nat)
  shows ∀ k, Nat(k) → choose(n, n + succ(k)) = 0
proof
  induction n
  | zero =>
    intros k kNat
    have shape : 0 + succ(k) = succ(k)
      exact add.left_unit at succ(k) with *
    rw shape
    exact choose.zero_left at k with *
  | succ(m) using ih =>
    intros k kNat
    have shape : succ(m) + succ(k) = succ(m + succ(k)) by ring
    rw [shape, choose.succ_succ]
    have first := ih with k, kNat
    have next : succ(m + succ(k)) = m + succ(succ(k)) by ring
    rw next
    have nextNat : Nat(succ(k)) by typecheck
    have second := ih with succ(k), nextNat
    rw [first, second]
    compute

theorem choose.diagonal (n : Nat)
  shows choose(n, n) = 1
proof
  induction n
  | zero => compute
  | succ(k) using ih =>
    rw [choose.succ_succ, ih]
    have zeroNat : Nat(0) by typecheck
    have above := choose.above at k with *
    have vanish := above with 0, zeroNat
    have shape : k + succ(0) = succ(k) by ring
    rw shape at vanish
    rw vanish
    compute

-- One row of the factorial form: from row m to row succ(m), by induction on
-- the column.
theorem choose.factorial_row (m k : Nat)
  requires row : ∀ i, Nat(i) → choose(m + i, i) * (i! * m!) = (m + i)!
  shows choose(succ(m) + k, k) * (k! * succ(m)!) = (succ(m) + k)!
proof
  induction k
  | zero =>
    have shape : succ(m) + 0 = succ(m) by ring
    rw shape
    rw choose.zero_right
    have zero : 0! = 1 by compute
    rw zero
    ring
  | succ(j) using inner =>
    have shape : succ(m) + succ(j) = succ(succ(m) + j) by ring
    rw shape
    rw choose.succ_succ
    have swap : succ(m) + j = m + succ(j) by ring
    have right := row with succ(j), assumption
    have factorialSucc : ∀ t, Nat(t) → succ(t)! = succ(t) * t!
      intros t tNat
      settle with factorial
    have expandJ := factorialSucc with j, assumption
    have expandM := factorialSucc with m, assumption
    have expandTop := factorialSucc with succ(m) + j, assumption
    rw expandTop
    rw expandJ
    rw expandM
    rw expandM at inner
    rw expandJ at right
    rw ← swap at right
    have regroup : (choose(succ(m) + j, j) + choose(succ(m) + j, succ(j))) * ((succ(j) * j!) * (succ(m) * m!)) = succ(j) * (choose(succ(m) + j, j) * (j! * (succ(m) * m!))) + succ(m) * (choose(succ(m) + j, succ(j)) * ((succ(j) * j!) * m!))
      ring
    rw regroup
    rw inner
    rw right
    ring

-- Pascal's rule and the factorial form agree: choose(n + k, k) · k! · n! = (n + k)!.
theorem choose.factorial (n : Nat)
  shows ∀ k, Nat(k) → choose(n + k, k) * (k! * n!) = (n + k)!
proof
  induction n
  | zero =>
    intros k kNat
    have shape : 0 + k = k
      exact add.left_unit at k with *
    rw [shape, choose.diagonal]
    have zero : 0! = 1 by compute
    rw zero
    ring
  | succ(m) using ih =>
    intros k kNat
    exact choose.factorial_row at m, k with *