-- Stirling numbers of the second kind by their triangle recurrence
-- S(n + 1, k + 1) = S(n, k) + (k + 1) · S(n, k + 1): the first and last
-- columns, vanishing above the diagonal, S(n + 1, 2) = 2 ^ n − 1,
-- S(n + 2, n + 1) = (n + 2 choose 2), positivity, and the Bell numbers
-- as row sums
-- 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.Stirling
import BinomialMore

-- Row 0 of the Stirling triangle: 1 at k = 0, then 0.
def stirlingZero (k : Nat) → Nat
  | zero => 1
  | succ(k) => 0

-- The recurrence at column k from the entries above and above-left: the
-- entry at column 0 is 0, and the entry at column succ(j) is the entry at
-- column j of the row above plus succ(j) times the entry at column succ(j).
def stirlingStep (k above left : Nat) → Nat
  | zero, above, left => 0
  | succ(j), above, left => add(left, mul(succ(j), above))

-- The number of partitions of an n-element set into k blocks.
def stirling (n k : Nat) → Nat
  | zero, k => stirlingZero(k)
  | succ(n), k => stirlingStep(k, stirling(n, k), stirling(n, predecessor(k)))
  descends n

theorem stirling.zero_zero
  shows stirling(0, 0) = 1
proof
  compute

theorem stirling.zero_succ (k : Nat)
  shows stirling(0, k + 1) = 0
proof
  have hk : k + 1 = succ(k) := add.right_one with *
  rw hk
  simp [stirling, stirlingZero]

theorem stirling.succ_zero (n : Nat)
  shows stirling(n + 1, 0) = 0
proof
  have hn : n + 1 = succ(n) := add.right_one with *
  rw hn
  simp [stirling, stirlingStep]

theorem stirling.succ_succ (n k : Nat)
  shows stirling(n + 1, k + 1) = stirling(n, k) + (k + 1) * stirling(n, k + 1)
proof
  have hn : n + 1 = succ(n) := add.right_one with *
  have hk : k + 1 = succ(k) := add.right_one with *
  rw [hn, hk]
  simp [stirling, stirlingStep, predecessor]

theorem stirling.natural (n : Nat)
  shows ∀ k, Nat(k) → Nat(stirling(n, k))
proof
  induction n
  | zero =>
    intros k kNat
    have split := natural.zero_or_successor at k
    cases split with isZero | isSucc
    . rw [isZero, stirling.zero_zero]
      typecheck
    . obtain ⟨j, jNat, shape⟩ := isSucc
      have hj : succ(j) = j + 1
        symm
        exact add.right_one at j
      rw [shape, hj, stirling.zero_succ]
      typecheck
  | m + 1 using ih =>
    intros k kNat
    have split := natural.zero_or_successor at k
    cases split with isZero | isSucc
    . rw [isZero, stirling.succ_zero]
      typecheck
    . obtain ⟨j, jNat, shape⟩ := isSucc
      have hj : succ(j) = j + 1
        symm
        exact add.right_one at j
      rw [shape, hj, stirling.succ_succ]
      have first := ih with j, jNat
      have nextNat : Nat(j + 1) by typecheck
      have second := ih with j + 1, nextNat
      typecheck

theorem stirling.natural_at (n k : Nat)
  shows Nat(stirling(n, k))
proof
  have all := stirling.natural at n
  exact all with k, assumption

theorem stirling.three_two
  shows stirling(3, 2) = 3
proof
  compute

theorem stirling.four_two
  shows stirling(4, 2) = 7
proof
  compute

theorem stirling.four_three
  shows stirling(4, 3) = 6
proof
  compute

theorem stirling.five_two
  shows stirling(5, 2) = 15
proof
  compute

theorem stirling.five_three
  shows stirling(5, 3) = 25
proof
  compute

theorem stirling.five_four
  shows stirling(5, 4) = 10
proof
  compute

-- One block: S(n + 1, 1) = 1.
theorem stirling.one (n : Nat)
  shows stirling(n + 1, 1) = 1
proof
  induction n
  | zero => compute
  | k + 1 using ih =>
    have step := stirling.succ_succ at k + 1, 0
    have one : 0 + 1 = 1 by compute
    rw one at step
    rw [step, stirling.succ_zero, ih]
    compute

-- Nothing above the diagonal.
theorem stirling.above (n : Nat)
  shows ∀ d, Nat(d) → stirling(n, n + 1 + d) = 0
proof
  induction n
  | zero =>
    intros d dNat
    have shape : 0 + 1 + d = d + 1 by ring
    rw shape
    exact stirling.zero_succ at d
  | k + 1 using ih =>
    intros d dNat
    have shape : k + 1 + 1 + d = (k + 1 + d) + 1 by ring
    rw [shape, stirling.succ_succ]
    have first := ih with d, dNat
    have nextNat : Nat(d + 1) by typecheck
    have second := ih with d + 1, nextNat
    have shape2 : k + 1 + (d + 1) = k + 1 + d + 1 by ring
    rw shape2 at second
    rw [first, second]
    ring

theorem stirling.above_diagonal (n k : Nat)
  requires below : n < k
  shows stirling(n, k) = 0
proof
  obtain ⟨d, dNat, shape⟩ := below
  have form : k = n + 1 + d
    rw ← shape
    ring
  rw form
  have all := stirling.above at n
  exact all with d, dNat

-- Singletons: S(n, n) = 1.
theorem stirling.diagonal (n : Nat)
  shows stirling(n, n) = 1
proof
  induction n
  | zero => compute
  | k + 1 using ih =>
    rw [stirling.succ_succ, ih]
    have below : k < k + 1 by arith
    have vanish := stirling.above_diagonal with below
    rw vanish
    ring

-- Two blocks: S(n + 1, 2) = 2 ^ n − 1.
theorem stirling.two (n : Nat)
  shows stirling(n + 1, 2) + 1 = 2 ^ n
proof
  induction n
  | zero => compute
  | k + 1 using ih =>
    have step := stirling.succ_succ at k + 1, 1
    have two : 1 + 1 = 2 by compute
    rw two at step
    rw [step, stirling.one, pow.add_one]
    have sNat := stirling.natural_at at k + 1, 2
    have powNat := pow.closed at 2, k
    ring [ih]

-- n + 1 blocks of n + 2 points: S(n + 2, n + 1) = (n + 2 choose 2).
theorem stirling.sub_one (n : Nat)
  shows 2 * stirling(n + 2, n + 1) = (n + 2) * (n + 1)
proof
  induction n
  | zero => compute
  | k + 1 using ih =>
    have shape : k + 1 + 2 = (k + 2) + 1 by ring
    have shape2 : k + 1 + 1 = k + 2 by ring
    rw [shape, stirling.succ_succ, shape2, stirling.diagonal]
    have sNat := stirling.natural_at at k + 2, k + 1
    ring [ih]

theorem stirling.sub_one_choose (n : Nat)
  shows stirling(n + 2, n + 1) = (n + 2) choose 2
proof
  have first := stirling.sub_one at n
  have second := choose.two at n + 1
  have shape : n + 1 + 1 = n + 2 by ring
  rw shape at second
  rw ← second at first
  have sNat := stirling.natural_at at n + 2, n + 1
  have cNat := choose.natural at n + 2, 2
  have twoPositive : 0 < 2 by arith
  have two := Less.ne_zero_of_positive with twoPositive
  exact mul.left_cancel_positive with two, first

-- Positivity on and below the diagonal.
theorem stirling.positive (n : Nat)
  shows ∀ k, Nat(k) → k ≤ n → 0 < stirling(n + 1, k + 1)
proof
  induction n
  | zero =>
    intros k kNat le
    have split := natural.zero_or_successor at k
    cases split with isZero | isSucc
    . rw isZero
      have one : 0 + 1 = 1 by compute
      have value : stirling(1, 1) = 1 by compute
      rw [one, value]
      arith
    . obtain ⟨i, iNat, shape⟩ := isSucc
      exfalso
      rw shape at le
      arith
  | j + 1 using ih =>
    intros k kNat le
    have split := natural.zero_or_successor at k
    cases split with isZero | isSucc
    . rw isZero
      have one : 0 + 1 = 1 by compute
      rw [one, stirling.one]
      arith
    . obtain ⟨i, iNat, shape⟩ := isSucc
      have hi : succ(i) = i + 1
        symm
        exact add.right_one at i
      rw [shape, hi, stirling.succ_succ]
      rw [shape, hi] at le
      have earlier : i ≤ j by arith
      have inner := ih with i, iNat, earlier
      have restNat := stirling.natural_at at j + 1, i + 1 + 1
      have productNat : Nat((i + 1 + 1) * stirling(j + 1, i + 1 + 1))
        typecheck
      arith

theorem stirling.positive_at (n k : Nat)
  requires le : k ≤ n
  shows 0 < stirling(n + 1, k + 1)
proof
  have all := stirling.positive at n
  exact all with k, assumption, le

-- Partial row sums of the Stirling triangle.
def bellSum (n k : Nat) → Nat
  | n, zero => 0
  | n, succ(k) => bellSum(n, k) + stirling(n, k)
  descends k

-- The Bell numbers: the number of partitions of an n-element set.
def bell (n : Nat) : Nat := bellSum(n, n + 1)

theorem bellSum.zero (n : Nat)
  shows bellSum(n, 0) = 0
proof
  simp [bellSum]

theorem bellSum.succ (n k : Nat)
  shows bellSum(n, k + 1) = bellSum(n, k) + stirling(n, k)
proof
  have hk : k + 1 = succ(k) := add.right_one with *
  rw hk
  simp [bellSum]

theorem bellSum.natural (n k : Nat)
  shows Nat(bellSum(n, k))
proof
  induction k
  | zero =>
    rw bellSum.zero
    typecheck
  | j + 1 using ih =>
    rw bellSum.succ
    have sNat := stirling.natural_at at n, j
    typecheck

theorem bell.value (n : Nat)
  shows bell(n) = bellSum(n, n + 1)
proof
  simp [bell]

theorem bell.natural (n : Nat)
  shows Nat(bell(n))
proof
  rw bell.value
  exact bellSum.natural at n, n + 1

theorem bell.zero
  shows bell(0) = 1
proof
  compute

theorem bell.one
  shows bell(1) = 1
proof
  compute

theorem bell.two
  shows bell(2) = 2
proof
  compute

theorem bell.three
  shows bell(3) = 5
proof
  compute

theorem bell.four
  shows bell(4) = 15
proof
  compute

theorem bell.five
  shows bell(5) = 52
proof
  compute

theorem bell.ge_one (n : Nat)
  shows 1 ≤ bell(n)
proof
  rw [bell.value, bellSum.succ, stirling.diagonal]
  have sumNat := bellSum.natural at n, n
  arith

-- The row sum is unchanged beyond the diagonal.
theorem bellSum.beyond (n d : Nat)
  shows bellSum(n, n + 1 + d) = bell(n)
proof
  induction d
  | zero =>
    have shape : n + 1 + 0 = n + 1 by ring
    rw [shape, bell.value]
    rfl
  | j + 1 using ih =>
    have shape : n + 1 + (j + 1) = (n + 1 + j) + 1 by ring
    rw [shape, bellSum.succ, ih]
    have below : n < n + 1 + j by arith
    have vanish := stirling.above_diagonal with below
    rw vanish
    have bNat := bell.natural at n
    ring