-- Vandermonde's identity: the convolution of two rows of Pascal's triangle
-- is an entry of a later row, the sum of the squares of a row is the central
-- binomial coefficient, and the bounds that follow from it
-- 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.Vandermonde
import Lucas
import Nat.SubtractionLemmas

-- The partial convolution ∑_{i < k} (m choose i) · (n choose (r − i)).
def chooseConvolution (m n r k : Nat) → Nat
  | m, n, r, zero => 0
  | m, n, r, succ(k) =>
    chooseConvolution(m, n, r, k) + choose(m, k) * choose(n, sub(r, k))
  descends k

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

theorem chooseConvolution.succ (m n r k : Nat)
  shows chooseConvolution(m, n, r, k + 1) =
    chooseConvolution(m, n, r, k) + (m choose k) * (n choose (r - k))
proof
  have hk : k + 1 = succ(k) := add.right_one with *
  rw hk
  simp [chooseConvolution]

theorem chooseConvolution.natural (m n r k : Nat)
  shows Nat(chooseConvolution(m, n, r, k))
proof
  induction k
  | zero =>
    rw chooseConvolution.zero
    typecheck
  | j + 1 using ih =>
    rw chooseConvolution.succ
    have leftNat := choose.natural at m, j
    have subNat := sub.closed at r, j
    have rightNat := choose.natural at n, r - j
    typecheck

theorem chooseConvolution.one (m n r : Nat)
  shows chooseConvolution(m, n, r, 1) = n choose r
proof
  have one : 0 + 1 = 1 by compute
  rw ← one
  rw [chooseConvolution.succ, chooseConvolution.zero, choose.zero_right]
  rw sub.zero_right
  have cNat := choose.natural at n, r
  ring

-- Row 0 of the triangle contributes only its first entry.
theorem chooseConvolution.zero_left (n r k : Nat)
  shows chooseConvolution(0, n, r, k + 1) = n choose r
proof
  induction k
  | zero =>
    have one : 0 + 1 = 1 by compute
    rw one
    exact chooseConvolution.one at 0, n, r
  | j + 1 using ih =>
    rw [chooseConvolution.succ, ih]
    have hj : j + 1 = succ(j) := add.right_one with *
    rw [hj, choose.zero_left]
    have cNat := choose.natural at n, r
    have subNat := sub.closed at r, succ(j)
    have rightNat := choose.natural at n, r - succ(j)
    ring

-- Pascal's rule for the convolution: the first row advances by one when the
-- target splits into the sums for r + 1 and for r.
theorem chooseConvolution.pascal (m n r k : Nat)
  shows chooseConvolution(m + 1, n, r + 1, k + 1) =
    chooseConvolution(m, n, r + 1, k + 1) + chooseConvolution(m, n, r, k)
proof
  induction k
  | zero =>
    have one : 0 + 1 = 1 by compute
    rw [one, chooseConvolution.one, chooseConvolution.one]
    rw chooseConvolution.zero
    have cNat := choose.natural at n, r + 1
    ring
  | j + 1 using ih =>
    have top := chooseConvolution.succ at m + 1, n, r + 1, j + 1
    have middle := chooseConvolution.succ at m, n, r + 1, j + 1
    have low := chooseConvolution.succ at m, n, r, j
    have pascalRule := choose.pascal at m, j
    have shift : r + 1 - (j + 1) = r - j := sub.add_one_both at r, j
    rw [top, ih, middle, low, pascalRule, shift]
    have aNat := chooseConvolution.natural at m, n, r + 1, j + 1
    have bNat := chooseConvolution.natural at m, n, r, j
    have cNat := choose.natural at m, j
    have dNat := choose.natural at m, j + 1
    have subNat := sub.closed at r, j
    have eNat := choose.natural at n, r - j
    ring

-- Vandermonde's identity ∑_{i ≤ r} (m choose i) · (n choose (r − i)) =
-- (m + n) choose r.
theorem choose.vandermonde (m n : Nat)
  shows ∀ r, Nat(r) → chooseConvolution(m, n, r, r + 1) = (m + n) choose r
proof
  induction m
  | zero =>
    intros r rNat
    have shape : 0 + n = n by ring
    rw [shape, chooseConvolution.zero_left]
    rfl
  | j + 1 using ih =>
    intros r rNat
    have split := natural.zero_or_successor at r
    cases split with isZero | isSucc
    . have one : 0 + 1 = 1 by compute
      rw [isZero, one, chooseConvolution.one]
      rw [choose.zero_right, choose.zero_right]
      rfl
    . obtain ⟨k, kNat, shape⟩ := isSucc
      have hk : succ(k) = k + 1
        symm
        exact add.right_one at k
      rw [shape, hk, chooseConvolution.pascal]
      have first := ih with k + 1, assumption
      have second := ih with k, kNat
      rw [first, second]
      have commute : j + 1 + n = (j + n) + 1 by ring
      rw [commute, choose.pascal]
      have aNat := choose.natural at j + n, k
      have bNat := choose.natural at j + n, k + 1
      ring

theorem choose.vandermonde_at (m n r : Nat)
  shows chooseConvolution(m, n, r, r + 1) = (m + n) choose r
proof
  have all := choose.vandermonde at m, n
  exact all with r, assumption

-- The convolution is symmetric in its two rows.
theorem chooseConvolution.symmetric (m n r : Nat)
  shows chooseConvolution(m, n, r, r + 1) = chooseConvolution(n, m, r, r + 1)
proof
  rw [choose.vandermonde_at, choose.vandermonde_at]
  have commute : m + n = n + m by ring
  rw commute
  rfl

theorem choose.vandermonde_one (m n : Nat)
  shows chooseConvolution(m, n, 1, 2) = m + n
proof
  have two : 1 + 1 = 2 by compute
  rw ← two
  rw [choose.vandermonde_at, choose.one_right]
  rfl

-- Partial sums of the squares of a row.
def chooseSquareSum (n k : Nat) → Nat
  | n, zero => 0
  | n, succ(k) => chooseSquareSum(n, k) + choose(n, k) * choose(n, k)
  descends k

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

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

theorem chooseSquareSum.natural (n k : Nat)
  shows Nat(chooseSquareSum(n, k))
proof
  induction k
  | zero =>
    rw chooseSquareSum.zero
    typecheck
  | j + 1 using ih =>
    rw chooseSquareSum.succ
    have cNat := choose.natural at n, j
    typecheck

-- Below the diagonal the self-convolution of a row is the sum of its squares.
theorem chooseSquareSum.convolution (n k : Nat)
  requires le : k ≤ n + 1
  shows chooseConvolution(n, n, n, k) = chooseSquareSum(n, k)
proof
  induction k
  | zero =>
    intro le
    rw [chooseConvolution.zero, chooseSquareSum.zero]
    rfl
  | j + 1 using ih =>
    intro le
    have earlier : j ≤ n + 1 by arith
    have inner := ih with earlier
    rw [chooseConvolution.succ, chooseSquareSum.succ, inner]
    have jLe : j ≤ n by arith
    have mirror := choose.symmetric at n, j with jLe
    rw ← mirror
    rfl

-- ∑_{k ≤ n} (n choose k)² = 2n choose n.
theorem choose.central_square_sum (n : Nat)
  shows chooseSquareSum(n, n + 1) = (2 * n) choose n
proof
  have le : n + 1 ≤ n + 1 by arith
  have convolution := chooseSquareSum.convolution at n, n + 1 with le
  rw ← convolution
  rw choose.vandermonde_at
  have double : n + n = 2 * n by ring
  rw double
  rfl

theorem chooseSquareSum.monotone (n k d : Nat)
  shows chooseSquareSum(n, k) ≤ chooseSquareSum(n, k + d)
proof
  induction d
  | zero =>
    have shape : k + 0 = k by ring
    rw shape
    have sumNat := chooseSquareSum.natural at n, k
    exact LessEqual.reflexive at chooseSquareSum(n, k)
  | j + 1 using ih =>
    have shape : k + (j + 1) = (k + j) + 1 by ring
    rw [shape, chooseSquareSum.succ]
    have sumNat := chooseSquareSum.natural at n, k
    have sumNatJ := chooseSquareSum.natural at n, k + j
    have cNat := choose.natural at n, k + j
    have squareNat : Nat((n choose (k + j)) * (n choose (k + j))) by typecheck
    arith

theorem chooseSquareSum.term_le (n k m : Nat)
  requires below : k < m
  shows (n choose k) * (n choose k) ≤ chooseSquareSum(n, m)
proof
  induction m
  | zero =>
    intro below
    have cNat := choose.natural at n, k
    have squareNat : Nat((n choose k) * (n choose k)) by typecheck
    arith
  | j + 1 using ih =>
    intro below
    rw chooseSquareSum.succ
    have sumNat := chooseSquareSum.natural at n, j
    have cNat := choose.natural at n, k
    have cNatJ := choose.natural at n, j
    have squareNat : Nat((n choose k) * (n choose k)) by typecheck
    have squareNatJ : Nat((n choose j) * (n choose j)) by typecheck
    have alternatives : k = j ∨ k ≠ j by tauto
    cases alternatives with same | different
    . rw same
      arith
    . have hj : j + 1 = succ(j) := add.right_one with *
      rw hj at below
      have strict := Less.of_successor_upper_ne with below, different
      have bound := ih with strict
      arith

-- Every square of a row entry is bounded by the central coefficient.
theorem choose.square_le_central (n k : Nat)
  requires le : k ≤ n
  shows (n choose k) * (n choose k) ≤ (2 * n) choose n
proof
  rw ← choose.central_square_sum
  have below : k < n + 1 by arith
  exact chooseSquareSum.term_le with below

theorem choose.le_central_of_le (n k : Nat)
  requires le : k ≤ n
  shows n choose k ≤ (2 * n) choose n
proof
  have square := choose.square_le_central with le
  have cNat := choose.natural at n, k
  have self := LessEqual.mul_self at n choose k
  have centralNat := choose.natural at 2 * n, n
  exact LessEqual.transitive with self, square

-- The row sum is bounded by the sum of the squares.
theorem chooseRowSum.le_square_sum (n k : Nat)
  shows chooseRowSum(n, k) ≤ chooseSquareSum(n, k)
proof
  induction k
  | zero =>
    rw [chooseRowSum.zero, chooseSquareSum.zero]
    exact LessEqual.reflexive at 0
  | j + 1 using ih =>
    rw [chooseRowSum.succ, chooseSquareSum.succ]
    have cNat := choose.natural at n, j
    have self := LessEqual.mul_self at n choose j
    have rowNat := chooseRowSum.natural at n, j
    have squareNat := chooseSquareSum.natural at n, j
    exact LessEqual.add with ih, self

-- The central binomial coefficient dominates 2 ^ n.
theorem choose.central_ge_pow (n : Nat)
  shows 2 ^ n ≤ (2 * n) choose n
proof
  rw [← choose.central_square_sum, ← chooseRowSum.full]
  exact chooseRowSum.le_square_sum at n, n + 1