Require Import Pascal.

Require Import Ring.
Require Import ArithRing.
Require Import Arith.

Fixpoint poly_sum_loop (f: nat -> nat -> nat) (remaining j: nat): nat :=
  match remaining with
  | 0 => 0
  | S i => f i j + poly_sum_loop f i (S j)
  end.

Definition poly_sum (f: nat -> nat -> nat) (n: nat): nat :=
  poly_sum_loop f (S n) 0.

Definition binomial_term (a b: nat): nat -> nat -> nat :=
  fun i j => pascal i j * a^i * b^j.

Definition binomial (a b: nat): nat -> nat :=
  poly_sum (binomial_term a b).

Lemma binomial_terms_add: forall a b i j: nat,
  a * binomial_term a b i (S j) + b * binomial_term a b (S i) j = binomial_term a b (S i) (S j).
Proof.
  intros.
  have (a * binomial_term a b i (S j) + b * binomial_term a b (S i) j).
  thatis (a * (pascal i (S j) * a^i * b^(S j)) + b * (pascal (S i) j * a^(S i) * b^j)).
  rearrange (pascal i (S j) * (a * a^i) * b^(S j) + pascal (S i) j * a^(S i) * (b * b^j)).
  thatis (pascal i (S j) * a^(S i) * b^(S j) + pascal (S i) j * a^(S i) * b^(S j)).
  rearrange ((pascal (S i) j + pascal i (S j)) * a^(S i) * b^(S j)).
  thatis (pascal (S i) (S j) * a^(S i) * b^(S j)).
  thatis (binomial_term a b (S i) (S j)).
  done.
Qed.

(* We want to show that (a + b) * binomial a b n = binomial a b (S n)
   but to build this up inductively, we need to look at the partial sum
   poly_sum_from i j (binomial_term a b) instead. Let's look at i=1.

       poly_sum_from 2 (1+j) (binomial_term a b)
     = binomial_term a b 1 (1+j) + binomial_term a b 0 (2+j)
     = pascal 1 (1+j)*a*b^(1+j) + pascal 0 (2+j)*b^(2+j)
     = (pascal 0 (1+j) + pascal 1 j)*a*b^(1+j) + pascal 0 (1+j)*b^(2+j)
     = a*(pascal 0 (1+j)*b^(1+j)) + b*(pascal 1 j*a*b^j + pascal 0 (1+j)*b^(1+j))
     = a*(poly_sum_from 1 (1+j) (binomial_term a b)) + b*(poly_sum_from 2 j (binomial_term a b))

   So this should be our inductive approach. *)

Lemma partial_binomial_equal: forall a b i j: nat,
  poly_sum_loop (binomial_term a b) (S i) (S j)
  = a*(poly_sum_loop (binomial_term a b) i (S j))
    + b*(poly_sum_loop (binomial_term a b) (S i) j).
Proof.
  intros a b i.
  set (f := binomial_term a b).
  set (loop := poly_sum_loop f).
  induction i; intro j.
  *both_sides_equal (b * b^j).
   -have (loop 1 (S j)).
    thatis (1*1*b^(S j) + 0).
    rearrange (b^(S j)).
    thatis (b * b^j).
    done.
   -have (a * loop 0 (S j) + b * loop 1 j).
    thatis (a * 0 + b * (1*1*b^j + 0)).
    rearrange (b * b^j).
    done.
  *have (loop (S (S i)) (S j)).
   thatis (binomial_term a b (S i) (S j) + loop (S i) (S (S j))).
   rewrite <- (binomial_terms_add a b i j).
   have (a * binomial_term a b i (S j) + b * binomial_term a b (S i) j + loop (S i) (S (S j))).
   thatis (a * f i (S j) + b * f (S i) j + loop (S i) (S (S j))).
   weknow Ind: (forall j: nat, loop (S i) (S j) = a * loop i (S j) + b * loop (S i) j).
   rewrite (Ind (S j)).
   have (a * f i (S j) + b * f (S i) j + (a * loop i (S (S j)) + b * loop (S i) (S j))).
   rearrange (a * (f i (S j) + loop i (S (S j))) + b * (f (S i) j + loop (S i) (S j))).
   thatis (a * (loop (S i) (S j)) + b * (loop (S (S i)) j)).
   done.
Qed.

Theorem binomial_step: forall a b n: nat,
  binomial a b (S n) = (a + b) * binomial a b n.
Proof.
  intros.
  unfold binomial.
  unfold poly_sum.
  set (loop := poly_sum_loop (binomial_term a b)).
  havegoal (loop (S (S n)) 0 = (a + b) * (loop (S n) 0)).
  (* The LHS naturally splits up into three terms, *)
  both_sides_equal (a * a^n + a * loop n 1 + b * loop (S n) 0).
  -have (loop (S (S n)) 0).
   thatis (binomial_term a b (S n) 0 + loop (S n) 1).
   thatis ((pascal (S n) 0 * (a * a^n) * 1) + loop (S n) 1).
   rewrite (pascal_0 (S n)); have ((1 * (a*a^n) * 1) + loop (S n) 1).
   rearrange (a * a^n + loop (S n) 1).
   assert (Eq: loop (S n) 1 = a * loop n 1 + b * loop (S n) 0). 
     { exact (partial_binomial_equal a b n 0). }
   rewrite Eq; have (a * a^n + (a * loop n 1 + b * loop (S n) 0)).
   rearrange (a * a^n + a * loop n 1 + b * loop (S n) 0).
   done.
  (* The RHS happens to also break down into these three terms. *)
  -have ((a + b) * loop (S n) 0).
   rearrange (a * loop (S n) 0 + b * loop (S n) 0).
   thatis (a * (pascal n 0 * a^n * 1 + loop n 1) + b * loop (S n) 0).
   rewrite (pascal_0 n); have (a * (1 * a^n * 1 + loop n 1) + b * loop (S n) 0).
   rearrange (a * a^n + a * loop n 1 + b * loop (S n) 0).
   done.
Qed.

Theorem binomial_expansion: forall a b n: nat,
  (a + b)^n = binomial a b n.
Proof.
  intros.
  induction n.
  *both_sides_equal 1.
   -have ((a + b)^0).
    thatis 1.
    done.
   -have (binomial a b 0).
    thatis (binomial_term a b 0 0 + 0).
    thatis (1 + 0).
    thatis 1.
    done.
  *have ((a + b)^(S n)).
   thatis ((a + b) * (a + b)^n).
   weknow Ind: ((a + b)^n = binomial a b n).
   rewrite Ind.
   havegoal ((a + b) * binomial a b n = binomial a b (S n)).
   exact (eq_sym (binomial_step a b n)).
Qed.