-- Trinomial coefficients (i + j + d)! / (i! j! d!) as products of binomial
-- coefficients: the factorial form, symmetry in the three parts, the
-- subset-of-a-subset identity, positivity, and the three-term recurrence
-- 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.Multinomial
import Stirling

-- The number of ways to split i + j + d points into parts of sizes i, j, d.
def trinomial (i j d : Nat) : Nat := choose(i + j + d, i) * choose(j + d, j)

theorem trinomial.value (i j d : Nat)
  shows trinomial(i, j, d) = ((i + j + d) choose i) * ((j + d) choose j)
proof
  simp [trinomial]

theorem trinomial.natural (i j d : Nat)
  shows Nat(trinomial(i, j, d))
proof
  rw trinomial.value
  have aNat := choose.natural at i + j + d, i
  have bNat := choose.natural at j + d, j
  typecheck

-- trinomial(i, j, d) · i! · j! · d! = (i + j + d)!.
theorem trinomial.factorial_form (i j d : Nat)
  shows trinomial(i, j, d) * (i! * j! * d!) = (i + j + d)!
proof
  rw trinomial.value
  have outerRow := choose.factorial at j + d
  have outer := outerRow with i, assumption
  have innerRow := choose.factorial at d
  have inner := innerRow with j, assumption
  have aNat := choose.natural at i + j + d, i
  have bNat := choose.natural at j + d, j
  have closedI := factorial.closed at i
  have closedJ := factorial.closed at j
  have closedD := factorial.closed at d
  have closedJD := factorial.closed at j + d
  ring [inner, outer]

theorem trinomial.symmetric (i j d : Nat)
  shows trinomial(i, j, d) = trinomial(j, i, d)
proof
  have first := trinomial.factorial_form at i, j, d
  have second := trinomial.factorial_form at j, i, d
  have aNat := trinomial.natural at i, j, d
  have bNat := trinomial.natural at j, i, d
  have closedI := factorial.closed at i
  have closedJ := factorial.closed at j
  have closedD := factorial.closed at d
  have balanced : (i! * j! * d!) * trinomial(i, j, d) =
      (i! * j! * d!) * trinomial(j, i, d)
    ring [first, second]
  have nonzeroI := factorial.nonzero at i
  have nonzeroJ := factorial.nonzero at j
  have nonzeroIJ := mul.nonzero with nonzeroI, nonzeroJ
  have productNat : Nat(i! * j!) by typecheck
  have nonzero := mul.nonzero with nonzeroIJ, (factorial.nonzero at d)
  exact mul.left_cancel_positive with nonzero, balanced

theorem trinomial.rotate (i j d : Nat)
  shows trinomial(i, j, d) = trinomial(j, d, i)
proof
  have first := trinomial.factorial_form at i, j, d
  have second := trinomial.factorial_form at j, d, i
  have aNat := trinomial.natural at i, j, d
  have bNat := trinomial.natural at j, d, i
  have closedI := factorial.closed at i
  have closedJ := factorial.closed at j
  have closedD := factorial.closed at d
  have balanced : (i! * j! * d!) * trinomial(i, j, d) =
      (i! * j! * d!) * trinomial(j, d, i)
    ring [first, second]
  have nonzeroI := factorial.nonzero at i
  have nonzeroJ := factorial.nonzero at j
  have nonzeroIJ := mul.nonzero with nonzeroI, nonzeroJ
  have productNat : Nat(i! * j!) by typecheck
  have nonzero := mul.nonzero with nonzeroIJ, (factorial.nonzero at d)
  exact mul.left_cancel_positive with nonzero, balanced

theorem trinomial.swap_last (i j d : Nat)
  shows trinomial(i, j, d) = trinomial(i, d, j)
proof
  have first := trinomial.symmetric at i, j, d
  have second := trinomial.rotate at j, i, d
  rw [first, second]
  rfl

-- Choosing i + j points and then i of them is choosing i points and then
-- j of the rest.
theorem choose.mul_choose (i j d : Nat)
  shows ((i + j + d) choose (i + j)) * ((i + j) choose i) = trinomial(i, j, d)
proof
  have form := trinomial.factorial_form at i, j, d
  have outerRow := choose.factorial at d
  have outer := outerRow with i + j, assumption
  have innerRow := choose.factorial at j
  have inner := innerRow with i, assumption
  have aNat := choose.natural at i + j + d, i + j
  have bNat := choose.natural at i + j, i
  have tNat := trinomial.natural at i, j, d
  have closedI := factorial.closed at i
  have closedJ := factorial.closed at j
  have closedD := factorial.closed at d
  have closedIJ := factorial.closed at i + j
  have balanced : (i! * j! * d!) *
      (((i + j + d) choose (i + j)) * ((i + j) choose i)) =
      (i! * j! * d!) * trinomial(i, j, d)
    ring [inner, outer, form]
  have nonzeroI := factorial.nonzero at i
  have nonzeroJ := factorial.nonzero at j
  have nonzeroIJ := mul.nonzero with nonzeroI, nonzeroJ
  have productNat : Nat(i! * j!) by typecheck
  have nonzero := mul.nonzero with nonzeroIJ, (factorial.nonzero at d)
  exact mul.left_cancel_positive with nonzero, balanced

theorem choose.subset_chain (i j d : Nat)
  shows ((i + j + d) choose (i + j)) * ((i + j) choose i) =
    ((i + j + d) choose i) * ((j + d) choose j)
proof
  rw [choose.mul_choose, trinomial.value]
  rfl

theorem trinomial.zero_right (i j : Nat)
  shows trinomial(i, j, 0) = (i + j) choose i
proof
  rw trinomial.value
  have shapeA : i + j + 0 = i + j by ring
  have shapeB : j + 0 = j by ring
  rw [shapeA, shapeB, choose.diagonal]
  have cNat := choose.natural at i + j, i
  ring

theorem trinomial.zero_left (j d : Nat)
  shows trinomial(0, j, d) = (j + d) choose j
proof
  rw trinomial.value
  have shape : 0 + j + d = j + d by ring
  rw [shape, choose.zero_right]
  have cNat := choose.natural at j + d, j
  ring

theorem trinomial.positive (i j d : Nat)
  shows 0 < trinomial(i, j, d)
proof
  rw trinomial.value
  have leI : i ≤ i + j + d by arith
  have leJ : j ≤ j + d by arith
  have first := choose.positive with leI
  have second := choose.positive with leJ
  have aNat := choose.natural at i + j + d, i
  have bNat := choose.natural at j + d, j
  exact Less.mul_positive with first, second

-- Pascal's rule for three parts.
theorem trinomial.recurrence (i j d : Nat)
  shows trinomial(i + 1, j + 1, d + 1) =
    trinomial(i, j + 1, d + 1) + trinomial(i + 1, j, d + 1) +
    trinomial(i + 1, j + 1, d)
proof
  rw [trinomial.value, trinomial.value, trinomial.value, trinomial.value]
  have shapeA : i + 1 + (j + 1) + (d + 1) = (i + j + d + 2) + 1 by ring
  have shapeB : i + (j + 1) + (d + 1) = i + j + d + 2 by ring
  have shapeC : i + 1 + j + (d + 1) = i + j + d + 2 by ring
  have shapeD : i + 1 + (j + 1) + d = i + j + d + 2 by ring
  have shapeE : j + 1 + (d + 1) = (j + d + 1) + 1 by ring
  have shapeF : j + (d + 1) = j + d + 1 by ring
  have shapeG : j + 1 + d = j + d + 1 by ring
  rw [shapeA, shapeB, shapeC, shapeD, shapeE, shapeF, shapeG]
  have outer := choose.pascal at i + j + d + 2, i
  have inner := choose.pascal at j + d + 1, j
  rw [outer, inner]
  have aNat := choose.natural at i + j + d + 2, i
  have bNat := choose.natural at i + j + d + 2, i + 1
  have cNat := choose.natural at j + d + 1, j
  have dNat := choose.natural at j + d + 1, j + 1
  ring