-- The binomial theorem: (A + B)^n is the sum over k ≤ n of
-- choose(n, k)·A^k·B^(n−k), by induction on n with Pascal's rule
-- 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.Real.BinomialTheorem
import BinomialTerms
import FiniteSum

theorem binomialTerm.sequence (A B : Real) (n : Nat)
  shows RealSequence(binomialTerm(A, B, n))
proof
  unfold RealSequence
  exact binomialTerm.function

-- Splitting the first term off a partial sum.
theorem realPartialSum.shift_one (u : Set) (m : Nat)
  requires sequence : RealSequence(u)
  shows ∑ i < succ(m), u(i) = (u(0):R) + (∑ k < m, shiftedRealSequence(u, succ(0))(k):R)
proof
  have shifted := shiftedRealSequence.sequence at u, succ(0) with sequence
  have split := realPartialSum.split_at at u, succ(0), m with sequence
  have form : add(succ(0), m) = succ(m) by ring
  rw form at split
  have one := realPartialSum.succ at u, 0
  have zero := realPartialSum.zero at u
  have firstReal := RealSequence.value at u, 0 with sequence
  have first : realPartialSum(u, succ(0)) = u(0)
    rw one
    linarith
  rw [split, first]
  rfl

-- The expansion by induction on n, with the bound written as succ(n) so that
-- the induction cases unfold the partial sum directly.
theorem binomial.expansion (A B : Real) (n : Nat)
  shows (A + B) ^ n = ∑ i < succ(n), binomialTerm(A, B, n)(i)
proof
  induction n
  | zero =>
    have sequence := binomialTerm.sequence at A, B, 0
    have zeroIndex := binomialTerm.zero_index at A, B, 0
    have shift := realPartialSum.shift_one at binomialTerm(A, B, 0), 0 with sequence
    have empty := realPartialSum.zero at shiftedRealSequence(binomialTerm(A, B, 0), succ(0))
    have sumReal : (A + B) ∈ ℝ by typecheck
    rw [realPow.zero, shift, zeroIndex, realPow.zero, empty]
    ring
  | succ(m) using ih =>
    have termsM := binomialTerm.sequence at A, B, m
    have termsNext := binomialTerm.sequence at A, B, succ(m)
    have shiftedNext := shiftedRealSequence.sequence at binomialTerm(A, B, succ(m)), succ(0) with termsNext
    have shiftedM := shiftedRealSequence.sequence at binomialTerm(A, B, m), succ(0) with termsM
    have constantA := constantSequence.sequence at A
    have constantB := constantSequence.sequence at B
    have scaledM := mulSequence.sequence with constantA, termsM
    have scaledShift := mulSequence.sequence with constantB, shiftedM
    have combined := addSequence.sequence with scaledM, scaledShift
    have agree : ∀ k ∈ ℕ, k < succ(m) → shiftedRealSequence(binomialTerm(A, B, succ(m)), succ(0))(k) = addSequence(mulSequence(constantSequence(A), binomialTerm(A, B, m)), mulSequence(constantSequence(B), shiftedRealSequence(binomialTerm(A, B, m), succ(0))))(k)
      intros k hk hkm
      have kNat : Nat(k) by typecheck
      have form : add(succ(0), k) = succ(k) by ring
      have shiftValue := shiftedRealSequence.value at binomialTerm(A, B, succ(m)), succ(0), k with termsNext
      have shiftValueM := shiftedRealSequence.value at binomialTerm(A, B, m), succ(0), k with termsM
      rw form at shiftValue
      rw form at shiftValueM
      have aReal := RealSequence.value at constantSequence(A), k with constantA
      have bReal := RealSequence.value at constantSequence(B), k with constantB
      have aValue := constantSequence.value at A, k
      have bValue := constantSequence.value at B, k
      have termReal := binomialTerm.real at A, B, m, k
      have shiftReal := RealSequence.value at shiftedRealSequence(binomialTerm(A, B, m), succ(0)), k with shiftedM
      have scaledValue := mulSequence.value at constantSequence(A), binomialTerm(A, B, m), k with constantA, termsM, aReal, termReal
      have scaledShiftValue := mulSequence.value at constantSequence(B), shiftedRealSequence(binomialTerm(A, B, m), succ(0)), k with constantB, shiftedM, bReal, shiftReal
      have scaledReal := RealSequence.value at mulSequence(constantSequence(A), binomialTerm(A, B, m)), k with scaledM
      have scaledShiftReal := RealSequence.value at mulSequence(constantSequence(B), shiftedRealSequence(binomialTerm(A, B, m), succ(0))), k with scaledShift
      have sumValue := addSequence.value at mulSequence(constantSequence(A), binomialTerm(A, B, m)), mulSequence(constantSequence(B), shiftedRealSequence(binomialTerm(A, B, m), succ(0))), k with scaledM, scaledShift, scaledReal, scaledShiftReal
      have pascal := binomialTerm.pascal at A, B, m, k
      rw [shiftValue, sumValue, scaledValue, scaledShiftValue, aValue, bValue, shiftValueM, pascal]
      rfl
    have equal := realPartialSum.eq_prefix at shiftedRealSequence(binomialTerm(A, B, succ(m)), succ(0)), addSequence(mulSequence(constantSequence(A), binomialTerm(A, B, m)), mulSequence(constantSequence(B), shiftedRealSequence(binomialTerm(A, B, m), succ(0)))), succ(m) with shiftedNext, combined, agree
    have add := realPartialSum.add at mulSequence(constantSequence(A), binomialTerm(A, B, m)), mulSequence(constantSequence(B), shiftedRealSequence(binomialTerm(A, B, m), succ(0))), succ(m) with scaledM, scaledShift
    have scaleA := realPartialSum.scale at binomialTerm(A, B, m), succ(m), A with termsM
    have scaleB := realPartialSum.scale at shiftedRealSequence(binomialTerm(A, B, m), succ(0)), succ(m), B with shiftedM
    have shiftNext := realPartialSum.shift_one at binomialTerm(A, B, succ(m)), succ(m) with termsNext
    have shiftM := realPartialSum.shift_one at binomialTerm(A, B, m), succ(m) with termsM
    have stepM := realPartialSum.succ at binomialTerm(A, B, m), succ(m)
    have le : m ≤ m
      witness 0
      ring
    have beyond := binomialTerm.beyond at A, B, m, m with le
    have firstNext := binomialTerm.succ_zero at A, B, m
    rw [realPow.succ, ih, shiftNext, firstNext, equal, add, scaleA, scaleB]
    rw shiftM at stepM
    have shiftSum : (∑ k < succ(m), shiftedRealSequence(binomialTerm(A, B, m), succ(0))(k):R) = (∑ i < succ(m), binomialTerm(A, B, m)(i):R) - (binomialTerm(A, B, m)(0):R) by linarith
    rw shiftSum
    ring

-- The binomial theorem: (A + B)ⁿ = ∑ₖ₌₀ⁿ C(n, k) Aᵏ Bⁿ⁻ᵏ, the sum running
-- over the n + 1 indices below n + 1.
theorem binomial.theorem (A B : Real) (n : Nat)
  shows (A + B) ^ n = ∑ i < n + 1, binomialTerm(A, B, n)(i)
proof
  have bound : n + 1 = succ(n) := add.right_one at n
  rw bound
  exact binomial.expansion at A, B, n with *