-- The derangement formula: d(n) = ∑_{k ≤ n} (−1)^k (n choose k) (n − k)!,
-- the inclusion–exclusion count of the permutations with no fixed point,
-- proved in the integers from the one-step recurrence
-- d(n + 1) = (n + 1) d(n) + (−1)^(n + 1).
-- 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.DerangementFormula
import Derangements
import Int

-- The partial sums ∑_{k < j} (−1)^k (n choose k) (n − k)!: the term of
-- index k counts the permutations of n points fixing k chosen points,
-- with the sign of inclusion–exclusion.
def derangementSum (n j : Nat) → Integer
  | n, zero => integerZero
  | n, succ(k) => derangementSum(n, k) +
      integerPow(-1, k) * integerOfNat(mul(choose(n, k), factorial(sub(n, k))))

theorem derangementSum.integer (n j : Nat)
  shows Integer(derangementSum(n, j))
proof
  induction j
  | zero =>
    simp [derangementSum]
    exact integerZero.integer
  | succ(k) using ih =>
    simp [derangementSum]
    have minus : Integer(-1) by typecheck
    have power := integerPow.integer at -1, k with minus
    have term := integerOfNat.integer at (n choose k) * (n - k)!
    typecheck

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

theorem derangementSum.succ (n k : Nat)
  shows derangementSum(n, k + 1) = derangementSum(n, k) +
    integerPow(-1, k) * integerOfNat((n choose k) * (n - k)!)
proof
  have one : k + 1 = succ(k) := add.right_one with *
  rw one
  simp [derangementSum]

-- The term of index k grows by the factor n + 1 when n grows by one.
theorem choose.factorial_step (n k : Nat)
  requires le : k ≤ n
  shows ((n + 1) choose k) * (n + 1 - k)! = (n + 1) * ((n choose k) * (n - k)!)
proof
  obtain ⟨d, dNat, shape⟩ := le
  have form : n = k + d
    rw ← shape
    rfl
  have difference : n - k = d
    have commute : k + d = d + k by ring
    rw [form, commute]
    exact sub.add_right at d, k
  have differenceSucc : n + 1 - k = d + 1
    have commute : k + d + 1 = (d + 1) + k by ring
    rw [form, commute]
    exact sub.add_right at d + 1, k
  have leSucc : k ≤ n + 1 by arith
  have top := choose.factorial_form at n + 1, k with leSucc
  have bottom := choose.factorial_form at n, k with le
  rw [difference, differenceSucc]
  rw differenceSucc at top
  rw difference at bottom
  have chooseTop := choose.natural at n + 1, k
  have chooseBottom := choose.natural at n, k
  have closedK := factorial.closed at k
  have closedD := factorial.closed at d
  have closedD1 := factorial.closed at d + 1
  have closedN := factorial.closed at n
  have left : k! * (((n + 1) choose k) * (d + 1)!) = (n + 1)!
    have reorder : k! * (((n + 1) choose k) * (d + 1)!) =
      ((n + 1) choose k) * (k! * (d + 1)!) by ring
    rw [reorder, top]
    rfl
  have right : k! * ((n + 1) * ((n choose k) * d!)) = (n + 1)!
    have reorder : k! * ((n + 1) * ((n choose k) * d!)) =
      (n + 1) * ((n choose k) * (k! * d!)) by ring
    rw [reorder, bottom, factorial.succ]
    rfl
  have balanced : k! * (((n + 1) choose k) * (d + 1)!) =
      k! * ((n + 1) * ((n choose k) * d!))
    rw [left, right]
    rfl
  exact mul.left_cancel_positive with (factorial.nonzero at k), balanced

-- Up to the last term, the sums for n + 1 and for n differ by the factor
-- n + 1.
theorem derangementSum.scale (n j : Nat)
  shows j ≤ n + 1 →
    derangementSum(n + 1, j) = integerOfNat(n + 1) * derangementSum(n, j)
proof
  induction j
  | zero =>
    intro le
    rw [derangementSum.zero, derangementSum.zero]
    have top := integerOfNat.integer at n + 1
    ring
  | k + 1 using ih =>
    intro le
    have kLe : k ≤ n by arith
    have kLeSucc : k ≤ n + 1 by arith
    have inner := ih with kLeSucc
    rw [derangementSum.succ, derangementSum.succ, inner]
    have term := choose.factorial_step with kLe
    rw [term, integerOfNat.mul]
    have minus : Integer(-1) by typecheck
    have power := integerPow.integer at -1, k with minus
    have sumInt := derangementSum.integer at n, k
    have topInt := integerOfNat.integer at n + 1
    have termInt := integerOfNat.integer at (n choose k) * (n - k)!
    ring

-- The one-step recurrence d(n + 1) = (n + 1) d(n) + (−1)^(n + 1), in ℤ.
theorem derangement.succ_integer (n : Nat)
  shows integerOfNat(derangement(n + 1)) =
    integerOfNat(n + 1) * integerOfNat(derangement(n)) + integerPow(-1, n + 1)
proof
  have minus : Integer(-1) by typecheck
  have alternatives := parity.exhaustive at n
  cases alternatives with even | odd
  . have iff := Even.iff_two_mul at n
    obtain ⟨k, kNat, shape⟩ := iff.forward with even
    rw shape
    have step := derangement.odd_succ at k
    have cast : integerOfNat(derangement(2 * k + 1) + 1) =
        integerOfNat((2 * k + 1) * derangement(2 * k))
      rw step
      rfl
    rw [integerOfNat.add, integerOfNat.mul, integerOfNat.one] at cast
    have sign : integerPow(-1, 2 * k + 1) = -1
      have double : 2 * k + 1 = k + k + 1 by ring
      rw [double, integerPow.neg_double_add_one, integerPow.one_left]
      rfl
    rw sign
    have dInt := integerOfNat.integer at derangement(2 * k + 1)
    have unfold : integerOfNat(derangement(2 * k + 1)) =
        (integerOfNat(derangement(2 * k + 1)) + 1) + -1 by ring
    rw [unfold, cast]
    have topInt := integerOfNat.integer at 2 * k + 1
    have previousInt := integerOfNat.integer at derangement(2 * k)
    ring
  . have iff := Odd.iff_two_mul_add_one at n
    obtain ⟨k, kNat, shape⟩ := iff.forward with odd
    rw shape
    have step := derangement.even_succ at k
    have two : 2 * k + 1 + 1 = 2 * k + 2 by ring
    rw two
    have cast : integerOfNat(derangement(2 * k + 2)) =
        integerOfNat((2 * k + 2) * derangement(2 * k + 1) + 1)
      rw step
      rfl
    rw [integerOfNat.add, integerOfNat.mul, integerOfNat.one] at cast
    have sign : integerPow(-1, 2 * k + 2) = 1
      have double : 2 * k + 2 = (k + 1) + (k + 1) by ring
      rw [double, integerPow.neg_double, integerPow.one_left]
      rfl
    rw [sign, cast]
    rfl

-- The derangement formula.
theorem derangement.closed_form (n : Nat)
  shows integerOfNat(derangement(n)) = derangementSum(n, n + 1)
proof
  induction n
  | zero =>
    have shape := derangementSum.succ at 0, 0
    have term : (0 choose 0) * (0 - 0)! = 1 by compute
    rw [derangement.zero, integerOfNat.one, shape, term]
    rw [derangementSum.zero, integerPow.zero, integerOfNat.one]
    ring
  | k + 1 using ih =>
    have step := derangement.succ_integer at k
    have last : ((k + 1) choose (k + 1)) * (k + 1 - (k + 1))! = 1
      rw [choose.diagonal, sub.self]
      compute
    have scale := derangementSum.scale at k, k + 1
    have scaled := scale with (LessEqual.reflexive at k + 1)
    rw [derangementSum.succ, last, integerOfNat.one, scaled, step, ih]
    have minus : Integer(-1) by typecheck
    have power := integerPow.integer at -1, k + 1 with minus
    have sumInt := derangementSum.integer at k, k + 1
    have topInt := integerOfNat.integer at k + 1
    ring