sum_first_n.lean

import Mathlib.Tactic
import Mathlib.Data.Real.Basic
import Mathlib.Data.Vector.Zip


theorem reverse_get_kth (n: Nat) (k: Fin n) (v: Vector Nat n):
  v.get k = v.reverse.get (Fin.revPerm k) := by
    cases n with
    | zero => exact k.elim0
    | succ =>
      rw [Vector.get_eq_get, Vector.get_eq_get]
      simp_rw [Vector.toList_reverse]
      rw [Fin.revPerm]
      rw [← Option.some_inj, ← List.get?_eq_get, ← List.get?_eq_get, List.get?_reverse]
      cases k using Fin.induction with
      | h0 =>
        simp
        congr
      | hs =>
        congr
        simp
        rw [Nat.sub_add_eq]
        refine Eq.symm (Nat.sub_eq_of_eq_add ?succ.hs.e_a.h)
        refine Eq.symm (Nat.add_sub_cancel' ?succ.hs.e_a.h.h)
        refine Iff.mpr Nat.succ_le ?succ.hs.e_a.h.h.a
        simp
      simp
      rw [@Nat.lt_succ]
      simp


theorem reverse_vector_sum (n: Nat) (v: Vector Nat n):
  v.toList.sum = v.reverse.toList.sum := by
    rw [← @List.sum_reverse]
    rfl


theorem sum_of_vectors_equals_sum_of_vectors_by_index (n: Nat) (v1: Vector Nat n) (v2: Vector Nat n) :
  (v1.map₂ Nat.add v2).toList.sum = v1.toList.sum + v2.toList.sum := by
    rw [Vector.sum_add_sum_eq_sum_zipWith]
    rfl


theorem sum_of_repl_product (n: Nat) (x: Nat):
  (Vector.replicate n x).toList.sum = n * x := by
    induction n with
    | zero =>
      rw [Vector.toList_empty]
      rw [List.sum_nil]
      rw [Nat.mul_comm]
      apply rfl
    | succ i ih =>
      rw [Nat.succ_mul]
      rw [Vector.replicate_succ]
      rw [@Vector.toList_cons]
      rw [@List.sum_cons]
      rw [ih]
      rw [Nat.add_comm x (i * x)]


---------------------------------------------

def first_n_nat : (n: Nat) → Vector Nat n
| 0 => Vector.nil
| n + 1 => (n + 1) ::ᵥ first_n_nat n


theorem kth_of_first_n_nat_equals_n_minus_k (n: Nat) (k: Fin n):
  (first_n_nat n).get k = n - k := by
    induction n with
    | zero => exact k.elim0
    | succ i ih =>
      rw [first_n_nat]
      cases k using Fin.cases with
      | H0 => rfl
      | Hs =>
        rw [Vector.get_cons_succ]
        rw [ih]
        simp


theorem kth_of_first_n_nat_rev_equals_succ_k (n: Nat) (k: Fin n):
  (first_n_nat n).reverse.get k = k + 1 := by
    rw [reverse_get_kth]
    rw [Vector.reverse_reverse]
    rw [kth_of_first_n_nat_equals_n_minus_k]
    simp
    refine tsub_eq_of_eq_add_rev ?h
    rw [Nat.sub_add_cancel]
    rw [Nat.succ_le]
    exact Fin.prop k


theorem kth_of_first_n_nat_plus_rev_equals_succ_n (n: Nat) (k: Fin n):
  (first_n_nat n).get k + (first_n_nat n).reverse.get k = n + 1 := by
    rw [kth_of_first_n_nat_equals_n_minus_k]
    rw [kth_of_first_n_nat_rev_equals_succ_k]
    rw [Nat.add_left_comm]
    rw [@Mathlib.Tactic.RingNF.add_assoc_rev]
    simp


theorem triangles_make_rectangle (n: Nat):
  Vector.map₂ Nat.add (first_n_nat n) (first_n_nat n).reverse = Vector.replicate n (n+1) := by
    apply Vector.ext
    intro m
    rw [Vector.get_replicate]
    rw [@Vector.get_map₂]
    simp
    rw [kth_of_first_n_nat_plus_rev_equals_succ_n]


theorem two_sum_eq_n_succ_n (n: Nat) :
  (first_n_nat n).toList.sum + (first_n_nat n).reverse.toList.sum = n * (n + 1) := by
    rw [<-sum_of_vectors_equals_sum_of_vectors_by_index]
    rw [triangles_make_rectangle]
    rw [sum_of_repl_product]