-- The derangement numbers 1, 0, 1, 2, 9, 44, …: the recurrence
-- d(n + 2) = (n + 1)(d(n + 1) + d(n)), the one-step recurrences
-- d(2n + 1) + 1 = (2n + 1) d(2n) and d(2n + 2) = (2n + 2) d(2n + 1) + 1,
-- positivity, monotonicity, and the comparison with n!
-- 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.Derangements
import FallingBounds

-- derangementAux(n, i, a, b) is the term n places after index i of a
-- sequence t with t(i) = a, t(i + 1) = b, and
-- t(k + 2) = (k + 1)(t(k + 1) + t(k)).
def derangementAux (n i a b : Nat) → Nat
  | zero, i, a, b => a
  | succ(k), i, a, b => derangementAux(k, i + 1, b, (i + 1) * (a + b))
  descends n

def derangement (n : Nat) : Nat := derangementAux(n, 0, 1, 0)

theorem derangementAux.zero (i a b : Nat)
  shows derangementAux(0, i, a, b) = a
proof
  simp [derangementAux]

theorem derangementAux.succ (k i a b : Nat)
  shows derangementAux(k + 1, i, a, b) =
    derangementAux(k, i + 1, b, (i + 1) * (a + b))
proof
  have one : k + 1 = succ(k) := add.right_one with *
  rw one
  simp [derangementAux]

theorem derangementAux.one (i a b : Nat)
  shows derangementAux(1, i, a, b) = b
proof
  simp [derangementAux]

theorem derangementAux.two (i a b : Nat)
  shows derangementAux(2, i, a, b) = (i + 1) * (a + b)
proof
  simp [derangementAux]

theorem derangementAux.closed (n : Nat)
  shows ∀ i, Nat(i) → ∀ a, Nat(a) → ∀ b, Nat(b) →
    Nat(derangementAux(n, i, a, b))
proof
  induction n
  | zero =>
    intros i iNat a aNat b bNat
    rw derangementAux.zero
    exact aNat
  | k + 1 using ih =>
    intros i iNat a aNat b bNat
    rw derangementAux.succ
    have nextNat : Nat(i + 1) by typecheck
    have productNat : Nat((i + 1) * (a + b)) by typecheck
    exact ih with i + 1, nextNat, b, bNat, (i + 1) * (a + b), productNat

-- Every such sequence obeys the derangement recurrence.
theorem derangementAux.shift (n : Nat)
  shows ∀ i, Nat(i) → ∀ a, Nat(a) → ∀ b, Nat(b) →
    derangementAux(n + 2, i, a, b) =
    (i + n + 1) * (derangementAux(n + 1, i, a, b) + derangementAux(n, i, a, b))
proof
  induction n
  | zero =>
    intros i iNat a aNat b bNat
    have two : 0 + 2 = 2 by compute
    have one : 0 + 1 = 1 by compute
    rw [two, one, derangementAux.two, derangementAux.one, derangementAux.zero]
    have shape : i + 0 + 1 = i + 1 by ring
    rw shape
    ring
  | k + 1 using ih =>
    intros i iNat a aNat b bNat
    have shape : k + 1 + 2 = (k + 2) + 1 by ring
    have shape2 : k + 1 + 1 = k + 2 by ring
    rw [shape, shape2]
    have top := derangementAux.succ at k + 2, i, a, b
    have middle := derangementAux.succ at k + 1, i, a, b
    have last := derangementAux.succ at k, i, a, b
    rw shape2 at middle
    rw [top, middle, last]
    have iNat1 : Nat(i + 1) by typecheck
    have cNat : Nat((i + 1) * (a + b)) by typecheck
    have inner := ih with i + 1, iNat1, b, bNat, (i + 1) * (a + b), cNat
    rw inner
    have aux1 := derangementAux.closed at k + 1
    have bNat1 := aux1 with i + 1, iNat1, b, bNat, (i + 1) * (a + b), cNat
    have aux0 := derangementAux.closed at k
    have bNat0 := aux0 with i + 1, iNat1, b, bNat, (i + 1) * (a + b), cNat
    ring

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

theorem derangement.one
  shows derangement(1) = 0
proof
  compute

theorem derangement.two
  shows derangement(2) = 1
proof
  compute

theorem derangement.three
  shows derangement(3) = 2
proof
  compute

theorem derangement.four
  shows derangement(4) = 9
proof
  compute

theorem derangement.five
  shows derangement(5) = 44
proof
  compute

theorem derangement.natural (n : Nat)
  shows Nat(derangement(n))
proof
  unfold derangement
  have zeroNat : Nat(0) by typecheck
  have oneNat : Nat(1) by typecheck
  exact derangementAux.closed at n with 0, zeroNat, 1, oneNat, 0, zeroNat

theorem derangement.recurrence (n : Nat)
  shows derangement(n + 2) = (n + 1) * (derangement(n + 1) + derangement(n))
proof
  unfold derangement
  have zeroNat : Nat(0) by typecheck
  have oneNat : Nat(1) by typecheck
  have shift := derangementAux.shift at n with 0, zeroNat, 1, oneNat, 0, zeroNat
  have shape : 0 + n + 1 = n + 1 by ring
  rw shape at shift
  exact shift

-- The one-step recurrences d(n + 1) = (n + 1) d(n) + (−1)^(n + 1), split by
-- the parity of n so that both sides stay natural.
theorem derangement.alternating_pair (n : Nat)
  shows derangement(2 * n + 1) + 1 = (2 * n + 1) * derangement(2 * n) ∧
    derangement(2 * n + 2) = (2 * n + 2) * derangement(2 * n + 1) + 1
proof
  induction n
  | zero =>
    constructor
    . compute
    . compute
  | k + 1 using ih =>
    have odd := ih.left
    have even := ih.right
    have shapeA : 2 * (k + 1) + 1 = 2 * k + 3 by ring
    have shapeB : 2 * (k + 1) = 2 * k + 2 by ring
    have shapeC : 2 * (k + 1) + 2 = 2 * k + 4 by ring
    rw [shapeA, shapeC, shapeB]
    have natA := derangement.natural at 2 * k
    have natB := derangement.natural at 2 * k + 1
    have natC := derangement.natural at 2 * k + 2
    have natD := derangement.natural at 2 * k + 3
    have stepOdd : derangement(2 * k + 3) + 1 =
        (2 * k + 3) * derangement(2 * k + 2)
      have step := derangement.recurrence at 2 * k + 1
      have three : 2 * k + 1 + 2 = 2 * k + 3 by ring
      have two : 2 * k + 1 + 1 = 2 * k + 2 by ring
      rw [three, two] at step
      rw [step, even]
      ring
    constructor
    . exact stepOdd
    . have step := derangement.recurrence at 2 * k + 2
      have four : 2 * k + 2 + 2 = 2 * k + 4 by ring
      have three : 2 * k + 2 + 1 = 2 * k + 3 by ring
      rw [four, three] at step
      rw step
      have reversed : (2 * k + 3) * derangement(2 * k + 2) =
          derangement(2 * k + 3) + 1
        symm
        exact stepOdd
      ring [reversed]

theorem derangement.odd_succ (n : Nat)
  shows derangement(2 * n + 1) + 1 = (2 * n + 1) * derangement(2 * n)
proof
  have pair := derangement.alternating_pair at n
  exact pair.left

theorem derangement.even_succ (n : Nat)
  shows derangement(2 * n + 2) = (2 * n + 2) * derangement(2 * n + 1) + 1
proof
  have pair := derangement.alternating_pair at n
  exact pair.right

-- d(n + 1) ≤ (n + 1) d(n) + 1 and (n + 1) d(n) ≤ d(n + 1) + 1 for every n.
theorem derangement.succ_le (n : Nat)
  shows derangement(n + 1) ≤ (n + 1) * derangement(n) + 1
proof
  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 natA := derangement.natural at 2 * k
    have natB := derangement.natural at 2 * k + 1
    arith
  . 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 natB := derangement.natural at 2 * k + 1
    have natC := derangement.natural at 2 * k + 2
    arith

theorem derangement.succ_ge (n : Nat)
  shows (n + 1) * derangement(n) ≤ derangement(n + 1) + 1
proof
  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 natA := derangement.natural at 2 * k
    have natB := derangement.natural at 2 * k + 1
    arith
  . 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 natB := derangement.natural at 2 * k + 1
    have natC := derangement.natural at 2 * k + 2
    arith

theorem derangement.positive (n : Nat)
  shows 0 < derangement(n + 2)
proof
  induction n
  | zero =>
    have two : 0 + 2 = 2 by compute
    rw [two, derangement.two]
    arith
  | k + 1 using ih =>
    rw derangement.recurrence
    have shape : k + 1 + 1 = k + 2 by ring
    rw shape
    have natA := derangement.natural at k + 1
    have natC := derangement.natural at k + 2
    have coefficient : 0 < k + 2 by arith
    have sum : 0 < derangement(k + 2) + derangement(k + 1) by arith
    exact Less.mul_positive with coefficient, sum

theorem derangement.ge_one (n : Nat)
  shows 1 ≤ derangement(n + 2)
proof
  have positive := derangement.positive at n
  have natural := derangement.natural at n + 2
  arith

theorem derangement.monotone (n : Nat)
  shows derangement(n + 2) ≤ derangement(n + 3)
proof
  have step := derangement.recurrence at n + 1
  have shape : n + 1 + 2 = n + 3 by ring
  have shape2 : n + 1 + 1 = n + 2 by ring
  rw [shape, shape2] at step
  rw step
  have natA := derangement.natural at n + 1
  have natB := derangement.natural at n + 2
  have expand : (n + 2) * (derangement(n + 2) + derangement(n + 1)) =
    derangement(n + 2) + (n + 1) * derangement(n + 2) +
    (n + 2) * derangement(n + 1)
    ring
  rw expand
  have productA : Nat((n + 1) * derangement(n + 2)) by typecheck
  have productB : Nat((n + 2) * derangement(n + 1)) by typecheck
  arith

-- d(n) ≤ n!, proved together with its successor.
theorem derangement.le_factorial_pair (n : Nat)
  shows derangement(n) ≤ n! ∧ derangement(n + 1) ≤ (n + 1)!
proof
  induction n
  | zero =>
    have one : 0 + 1 = 1 by compute
    have zeroFactorial : 0! = 1 by compute
    rw [one, derangement.zero, derangement.one, zeroFactorial, factorial.one]
    constructor
    . arith
    . arith
  | k + 1 using ih =>
    have first := ih.left
    have second := ih.right
    constructor
    . exact second
    . have shape : k + 1 + 1 = k + 2 by ring
      rw [shape, derangement.recurrence]
      have natA := derangement.natural at k
      have natB := derangement.natural at k + 1
      have closedK := factorial.closed at k
      have closedK1 := factorial.closed at k + 1
      have sum := LessEqual.add with second, first
      have scaled := LessEqual.mul_left at derangement(k + 1) + derangement(k),
        (k + 1)! + k!, k + 1 with sum
      have expand : (k + 2)! = (k + 1) * ((k + 1)! + k!)
        have shape2 : k + 2 = k + 1 + 1 by ring
        rw [shape2, factorial.succ, factorial.succ]
        ring
      rw expand
      exact scaled

theorem derangement.le_factorial (n : Nat)
  shows derangement(n) ≤ n!
proof
  have pair := derangement.le_factorial_pair at n
  exact pair.left

-- d(n + 1) < (n + 1)!, proved together with its successor.
theorem derangement.lt_factorial_pair (n : Nat)
  shows derangement(n + 1) < (n + 1)! ∧ derangement(n + 2) < (n + 2)!
proof
  induction n
  | zero =>
    have one : 0 + 1 = 1 by compute
    have two : 0 + 2 = 2 by compute
    rw [one, two, derangement.one, derangement.two]
    rw [factorial.one, factorial.two]
    constructor
    . arith
    . arith
  | k + 1 using ih =>
    have first := ih.left
    have second := ih.right
    have shape : k + 1 + 1 = k + 2 by ring
    rw shape
    constructor
    . exact second
    . have step := derangement.recurrence at k + 1
      have shape2 : k + 1 + 2 = k + 3 by ring
      rw [shape2, shape] at step
      rw [shape2, step]
      have natA := derangement.natural at k + 1
      have natB := derangement.natural at k + 2
      have closedK1 := factorial.closed at k + 1
      have closedK2 := factorial.closed at k + 2
      have firstLe := derangement.le_factorial at k + 1
      have sum := Less.add_LessEqual with second, firstLe
      have coefficient : 0 < k + 2 by arith
      have nonzero := Less.ne_zero_of_positive with coefficient
      have scaled := Less.mul_left at derangement(k + 2) + derangement(k + 1),
        (k + 2)! + (k + 1)!, k + 2 with nonzero, sum
      have expand : (k + 3)! = (k + 2) * ((k + 2)! + (k + 1)!)
        have shape3 : k + 3 = k + 2 + 1 by ring
        rw [shape3, factorial.succ]
        have inner := factorial.succ at k + 1
        rw shape at inner
        rw inner
        ring
      rw expand
      exact scaled

theorem derangement.lt_factorial (n : Nat)
  shows derangement(n + 1) < (n + 1)!
proof
  have pair := derangement.lt_factorial_pair at n
  exact pair.left