-- ArkLib/Data/FieldTheory/AdditiveNTT/AdditiveNTT.lean

import ArkLib.Data.FieldTheory.AdditiveNTT.NovelPolynomialBasis

import Mathlib.Tactic

import Mathlib.Data.Finsupp.Defs

import Mathlib.LinearAlgebra.LinearIndependent.Defs

open Polynomial AdditiveNTT Module

namespace AdditiveNTT

variable {r : ℕ} [NeZero r]

variable {L : Type u} [Field L] [Fintype L] [DecidableEq L]

variable (𝔽q : Type u) [Field 𝔽q] [Fintype 𝔽q] [DecidableEq 𝔽q]
  [h_Fq_char_prime : Fact (Nat.Prime (ringChar 𝔽q))] [hF₂ : Fact (Fintype.card 𝔽q = 2)]

variable [Algebra 𝔽q L]

variable (β : Fin r → L) [hβ_lin_indep : Fact (LinearIndependent 𝔽q β)]
  [h_β₀_eq_1 : Fact (β 0 = 1)]

variable {ℓ R_rate : ℕ} (h_ℓ_add_R_rate : ℓ + R_rate < r)-- ℓ ∈ {1, ..., r-1}

section IntermediateStructures

noncomputable def qMap (i : Fin r) : L[X] :=
  let constMultiplier := ((W 𝔽q β i).eval (β i))^(Fintype.card 𝔽q)
    / ((W 𝔽q β (i + 1)).eval (β (i + 1)))
  C constMultiplier * ∏ c: 𝔽q, (X - C (algebraMap 𝔽q L c))

noncomputable def qCompositionChain (i : Fin r) : L[X] :=
  match i with
  | ⟨0, _⟩ => X
  | ⟨k + 1, h_k_add_1⟩ => (qMap 𝔽q β ⟨k, by admit /- proof elided -/
  ⟩).comp (qCompositionChain ⟨k, by admit /- proof elided -/
  ⟩)

noncomputable section DomainBijection

end DomainBijection

noncomputable def intermediateNormVpoly
    
    (i: Fin (ℓ+1)) (k : Fin (ℓ - i + 1)) : L[X] :=
  
  Fin.foldl (n:=k) (fun acc j =>
    (qMap 𝔽q β ⟨(i : ℕ) + (j : ℕ), by admit /- proof elided -/
    ⟩).comp acc) (X)

noncomputable def intermediateNovelBasisX (i : Fin (ℓ + 1)) (j : Fin (2 ^ (ℓ - i))) : L[X] :=
  (Finset.univ: Finset (Fin (ℓ - i)) ).prod (fun k =>
    (intermediateNormVpoly 𝔽q β h_ℓ_add_R_rate i (k:=⟨k, by admit /- proof elided -/
    ⟩)) ^ (Nat.getBit k j))

noncomputable def intermediateEvaluationPoly (i : Fin (ℓ + 1))
    (coeffs : Fin (2 ^ (ℓ - i)) → L) : L[X] :=
  ∑ (⟨j, hj⟩: Fin (2^(ℓ-i))), C (coeffs ⟨j, by admit /- proof elided -/
  ⟩) *
    (intermediateNovelBasisX 𝔽q β h_ℓ_add_R_rate i ⟨j, by admit /- proof elided -/
    ⟩)

end IntermediateStructures

section AlgorithmCorrectness

noncomputable def evaluationPointω (i : Fin (ℓ + 1))
    (x : Fin (2 ^ (ℓ + R_rate - i))) : L := 
    
  ∑ (⟨k, hk⟩: Fin (ℓ + R_rate - i)),
    if Nat.getBit k x.val = 1 then
      (normalizedW 𝔽q β ⟨i, by admit /- proof elided -/
      ⟩).eval (β ⟨i + k, by admit /- proof elided -/
      ⟩)
    else
      0

def tileCoeffs (a : Fin (2 ^ ℓ) → L) : Fin (2^(ℓ + R_rate)) → L :=
  fun v => a (Fin.mk (v.val % (2^ℓ)) (Nat.mod_lt v.val (pow_pos (zero_lt_two) ℓ)))

def coeffsBySuffix (a : Fin (2 ^ ℓ) → L) (i : Fin (ℓ + 1)) (v : Fin (2 ^ i.val)) :
  Fin (2 ^ (ℓ - i)) → L :=
  fun ⟨j, hj⟩ => by admit /- proof elided -/

def additiveNTTInvariant (evaluation_buffer : Fin (2 ^ (ℓ + R_rate)) → L)
    (original_coeffs : Fin (2 ^ ℓ) → L) (i : Fin (ℓ + 1)) : Prop :=
  ∀ (j : Fin (2^(ℓ + R_rate))),
    let u_b_v := j.val
    let v: Fin (2^i.val) := ⟨Nat.getLowBits i.val u_b_v, by admit /- proof elided -/
    ⟩ 
    let u_b := u_b_v / (2^i.val) 
    have h_u_b : u_b = u_b_v / (2^i.val) := by admit /- proof elided -/

-- Repository definitions used:
content: noncomputable def W (i : Fin r) : L[X] :=
  ∏ u : U 𝔽q β i, (X - C u.val)

content: def U (i : Fin r) : Subspace 𝔽q L := Submodule.span 𝔽q (β '' (Set.Ico 0 i))

content: noncomputable def normalizedW (i : Fin r) : L[X] :=
  C (1 / (W 𝔽q β i).eval (β i)) * W 𝔽q β i

content: def getBit (k n : Nat) : Nat := (n >>> k) &&& 1

content: def getLowBits (numLowBits : ℕ) (n : ℕ) := n &&& ((1 <<< numLowBits) - 1)

content: @[elab_as_elim] def Fin.succRecOnSameFinType {motive : Fin r → Sort _}
    (zero : motive (0 : Fin r))
    (succ : ∀ i : Fin r, i + 1 < r → motive i → motive (i + 1)) : ∀ (i : Fin r), motive i
  | ⟨0, _⟩ => by admit /- proof elided -/
  | ⟨Nat.succ i_val, h⟩ => by admit /- proof elided -/

content: @[simps!]
noncomputable def algEquivAevalXSubC {R : Type*} [CommRing R] (t : R) : R[X] ≃ₐ[R] R[X] :=

-- Repository lemmas used:
content: lemma Xⱼ_zero_eq_one (ℓ : ℕ) (h_ℓ : ℓ ≤ r) :
  Xⱼ 𝔽q β ℓ h_ℓ ⟨0, by exact Nat.two_pow_pos ℓ⟩ = 1

content: omit hF₂ in
theorem W_linear_comp_decomposition (i : Fin r) (h_i_add_1 : i + 1 < r) :
    ∀ p: L[X], (W 𝔽q β (i + 1)).comp p =
      ((W 𝔽q β i).comp p) ^ Fintype.card 𝔽q -
        C (eval (β i) (W 𝔽q β i)) ^ (Fintype.card 𝔽q - 1) * ((W 𝔽q β i).comp p)

content: theorem W_linearity (i : Fin r)
    : IsLinearMap 𝔽q (f := fun inner_p ↦ (W 𝔽q β i).comp inner_p)

content: omit hF₂ in
lemma inductive_linear_map_W (i : Fin r) (h_i_add_1 : i + 1 < r)
    (h_prev_linear_map : IsLinearMap 𝔽q (f := fun inner_p ↦ (W 𝔽q β i).comp inner_p))
    : IsLinearMap 𝔽q (f := fun inner_p ↦ (W 𝔽q β (i + 1)).comp inner_p)

content: omit h_Fq_char_prime hF₂ in
lemma inductive_rec_form_W_comp (i : Fin r) (h_i_add_1 : i + 1 < r)
    (h_prev_linear_map : IsLinearMap (R := 𝔽q) (M := L[X]) (M₂ := L[X])
      (f := fun inner_p ↦ (W 𝔽q β i).comp inner_p))
    : ∀ p: L[X], (W 𝔽q β (i + 1)).comp p =
      ((W 𝔽q β i).comp p) ^ Fintype.card 𝔽q -
        C (eval (β i) (W 𝔽q β i)) ^ (Fintype.card 𝔽q - 1) * ((W 𝔽q β i).comp p)

content: lemma Wᵢ_eval_βᵢ_neq_zero
    (i : Fin r): (W 𝔽q β i).eval (β i) ≠ 0

content: lemma βᵢ_not_in_Uᵢ (i : Fin r) :
    β i ∉ U 𝔽q β i

content: lemma W_prod_comp_decomposition
    (i : Fin r) (hi : i > 0) :
    (W 𝔽q β i) = ∏ c: 𝔽q, (W 𝔽q β (i-1)).comp (X - C (c • β (i-1)))

content: lemma W_splits (i : Fin r) : (W 𝔽q β i).Splits (RingHom.id L)

content: omit h_Fq_char_prime hF₂ in
lemma rootMultiplicity_prod_W_comp_X_sub_C
    (i : Fin r) (h_i_add_1 : i + 1 < r) (a : L) :
    rootMultiplicity a ((univ : Finset 𝔽q).prod (fun c => (W 𝔽q β i).comp (X - C (c • β i)))) =
    if a ∈ (U 𝔽q β (i+1) : Set L) then 1 else 0

content: omit [DecidableEq L] h_Fq_char_prime hF₂ hβ_lin_indep in
lemma Prod_W_comp_X_sub_C_ne_zero (i : Fin r) :
    (univ : Finset 𝔽q).prod (fun c => (W 𝔽q β i).comp (X - C (c • β i))) ≠ 0

content: omit [DecidableEq L] [Fintype 𝔽q] hβ_lin_indep in
lemma W_ne_zero (i : Fin r) : (W 𝔽q β i) ≠ 0

content: lemma roots_comp_X_sub_C (p : L[X]) (a : L) :
    (p.comp (X - C a)).roots = p.roots.map (fun r => r + a)

content: lemma rootMultiplicity_comp_X_sub_C (p : L[X]) (a x : L) :
    rootMultiplicity x (p.comp (X - C a)) = rootMultiplicity (x - a) p

content: omit [Fintype L] [DecidableEq L] in
lemma comp_X_sub_C_eq_zero_iff (p : L[X]) (a : L) :
  p.comp (X - C a) = 0 ↔ p = 0

content: lemma rootMultiplicity_W (i : Fin r) (a : L) :
    rootMultiplicity a (W 𝔽q β i) = if a ∈ (U 𝔽q β i : Set L) then 1 else 0

content: lemma roots_W (i : Fin r) : -- converts root Multiset into (univ: Uᵢ.val.map)
  (W 𝔽q β i).roots = (univ : Finset (U 𝔽q β i)).val.map (fun u => u.val)

content: omit [Fintype L] [DecidableEq L] [Fintype 𝔽q] h_Fq_char_prime hβ_lin_indep in
theorem root_U_lift_up (i : Fin r) (h_i_add_1 : i + 1 < r) (a : L) (x : 𝔽q) :
  a - x • β i ∈ (U 𝔽q β i) → a ∈ (U 𝔽q β (i+1))

content: omit [Fintype L] [DecidableEq L] [Fintype 𝔽q] h_Fq_char_prime in
theorem root_U_lift_down
  (i : Fin r) (h_i_add_1 : i + 1 < r) (a : L) :
  a ∈ (U 𝔽q β (i+1)) → ∃! x: 𝔽q, a - x • β i ∈ (U 𝔽q β i)

content: lemma W_monic (i : Fin r) : (W 𝔽q β i).Monic

content: lemma comp_sub_C_of_linear_eval (p : L[X])
  (h_lin : IsLinearMap 𝔽q (f := fun inner_p ↦ p.comp inner_p)) (a : L) :
    p.comp (X - C a) = p - C (eval a p)

content: theorem prod_poly_sub_C_eq_poly_pow_card_sub_poly_in_L
  (p : L[X]) :
  (∏ c ∈ (Finset.univ : Finset Fq), (p - Polynomial.C (algebraMap Fq L c))) =
    p^(Fintype.card Fq) - p

content: theorem prod_X_sub_C_eq_X_pow_card_sub_X_in_L :
  (∏ c ∈ (Finset.univ : Finset Fq), (Polynomial.X - Polynomial.C (algebraMap Fq L c))) =
    Polynomial.X^(Fintype.card Fq) - Polynomial.X

content: theorem prod_X_sub_C_eq_X_pow_card_sub_X :
  (∏ c ∈ (Finset.univ : Finset Fq), (Polynomial.X - Polynomial.C c)) =
    Polynomial.X^(Fintype.card Fq) - Polynomial.X

content: omit [DecidableEq L] [Fintype 𝔽q] hβ_lin_indep in
lemma W₀_eq_X : W 𝔽q β 0 = X

content: lemma Fin.val_add_one' (a : Fin r) (h_a_add_1 : a + 1 < r) : (a + 1).val = a.val + 1

content: lemma getLowBits_eq_mod_two_pow {numLowBits : ℕ} (n : ℕ) :
  getLowBits numLowBits n = n % (2 ^ numLowBits)

content: lemma getLowBits_lt_two_pow {n : ℕ} (numLowBits : ℕ) :
    getLowBits numLowBits n < 2 ^ numLowBits

-- Local definitions used:
content: noncomputable def qMap (i : Fin r) : L[X] :=
  let constMultiplier := ((W 𝔽q β i).eval (β i))^(Fintype.card 𝔽q)
    / ((W 𝔽q β (i + 1)).eval (β (i + 1)))
  C constMultiplier * ∏ c: 𝔽q, (X - C (algebraMap 𝔽q L c))

content: noncomputable def qCompositionChain (i : Fin r) : L[X] :=
  match i with
  | ⟨0, _⟩ => X
  | ⟨k + 1, h_k_add_1⟩ => (qMap 𝔽q β ⟨k, by admit /- proof elided -/
  ⟩).comp (qCompositionChain ⟨k, by admit /- proof elided -/
  ⟩)

content: noncomputable def intermediateNormVpoly
    
    (i: Fin (ℓ+1)) (k : Fin (ℓ - i + 1)) : L[X] :=
  
  Fin.foldl (n:=k) (fun acc j =>
    (qMap 𝔽q β ⟨(i : ℕ) + (j : ℕ), by admit /- proof elided -/
    ⟩).comp acc) (X)

content: noncomputable def intermediateNovelBasisX (i : Fin (ℓ + 1)) (j : Fin (2 ^ (ℓ - i))) : L[X] :=
  (Finset.univ: Finset (Fin (ℓ - i)) ).prod (fun k =>
    (intermediateNormVpoly 𝔽q β h_ℓ_add_R_rate i (k:=⟨k, by admit /- proof elided -/
    ⟩)) ^ (Nat.getBit k j))

content: noncomputable def intermediateEvaluationPoly (i : Fin (ℓ + 1))
    (coeffs : Fin (2 ^ (ℓ - i)) → L) : L[X] :=
  ∑ (⟨j, hj⟩: Fin (2^(ℓ-i))), C (coeffs ⟨j, by admit /- proof elided -/
  ⟩) *
    (intermediateNovelBasisX 𝔽q β h_ℓ_add_R_rate i ⟨j, by admit /- proof elided -/
    ⟩)

content: noncomputable def evaluationPointω (i : Fin (ℓ + 1))
    (x : Fin (2 ^ (ℓ + R_rate - i))) : L := 
    
  ∑ (⟨k, hk⟩: Fin (ℓ + R_rate - i)),
    if Nat.getBit k x.val = 1 then
      (normalizedW 𝔽q β ⟨i, by admit /- proof elided -/
      ⟩).eval (β ⟨i + k, by admit /- proof elided -/
      ⟩)
    else
      0

content: def tileCoeffs (a : Fin (2 ^ ℓ) → L) : Fin (2^(ℓ + R_rate)) → L :=
  fun v => a (Fin.mk (v.val % (2^ℓ)) (Nat.mod_lt v.val (pow_pos (zero_lt_two) ℓ)))

content: def coeffsBySuffix (a : Fin (2 ^ ℓ) → L) (i : Fin (ℓ + 1)) (v : Fin (2 ^ i.val)) :
  Fin (2 ^ (ℓ - i)) → L :=
  fun ⟨j, hj⟩ => by admit /- proof elided -/

content: def additiveNTTInvariant (evaluation_buffer : Fin (2 ^ (ℓ + R_rate)) → L)
    (original_coeffs : Fin (2 ^ ℓ) → L) (i : Fin (ℓ + 1)) : Prop :=
  ∀ (j : Fin (2^(ℓ + R_rate))),
    let u_b_v := j.val
    let v: Fin (2^i.val) := ⟨Nat.getLowBits i.val u_b_v, by admit /- proof elided -/
    ⟩ 
    let u_b := u_b_v / (2^i.val) 
    have h_u_b : u_b = u_b_v / (2^i.val) := by admit /- proof elided -/

-- Local lemmas used:
content: lemma qMap_comp_normalizedW (i : Fin r) (h_i_add_1 : i + 1 < r) :
  (qMap 𝔽q β i).comp (normalizedW 𝔽q β i) = normalizedW 𝔽q β (i + 1)

content: lemma qCompositionChain_eq_foldl (i : Fin r) :
  qCompositionChain 𝔽q β (ℓ:=ℓ) (R_rate:=R_rate) i =
  Fin.foldl (n:=i) (fun acc j =>
    (qMap 𝔽q β ⟨j, by omega⟩).comp acc) (X)

content: lemma normalizedW_eq_qMap_composition (ℓ R_rate : ℕ) (i : Fin r) :
    normalizedW 𝔽q β i = qCompositionChain 𝔽q β (ℓ:=ℓ) (R_rate:=R_rate) i

content: theorem base_intermediateNormVpoly
  (k : Fin (ℓ + 1)) :
  intermediateNormVpoly 𝔽q β h_ℓ_add_R_rate ⟨0, by
    by_contra ht
    simp only [not_lt, nonpos_iff_eq_zero] at ht
    contradiction
  ⟩ ⟨k, by simp only [tsub_zero]; omega⟩ =
  normalizedW 𝔽q β ⟨k, by omega⟩

content: theorem base_intermediateNovelBasisX (j : Fin (2 ^ ℓ)) :
  intermediateNovelBasisX 𝔽q β h_ℓ_add_R_rate ⟨0, by
    by_contra ht
    simp only [not_lt, nonpos_iff_eq_zero] at ht
    contradiction
  ⟩ j =
  Xⱼ 𝔽q β ℓ (by omega) j

-- Target theorem and proof:
omit [DecidableEq 𝔽q] hF₂ in
lemma initial_tiled_coeffs_correctness (h_ℓ : ℓ ≤ r) (a : Fin (2 ^ ℓ) → L) :
    let b: Fin (2^(ℓ + R_rate)) → L := tileCoeffs a
    additiveNTTInvariant 𝔽q β h_ℓ_add_R_rate b a (i := ⟨ℓ, by omega⟩)
:= tileCoeffs a
    additiveNTTInvariant 𝔽q β h_ℓ_add_R_rate b a (i := ⟨ℓ, by omega⟩) := by
    unfold additiveNTTInvariant
    simp only
    intro j
    unfold coeffsBySuffix
    simp only [tileCoeffs, evaluationPointω, intermediateEvaluationPoly, Fin.eta]
    have h_ℓ_sub_ℓ: 2^(ℓ - ℓ) = 1 := by norm_num

    set f_right: Fin (2^(ℓ - ℓ)) → L[X] :=
      fun ⟨x, hx⟩ => C (a ⟨↑x <<< ℓ ||| Nat.getLowBits ℓ (↑j), by
        simp only [tsub_self, pow_zero, Nat.lt_one_iff] at hx
        simp only [hx, Nat.zero_shiftLeft, Nat.zero_or]
        exact Nat.getLowBits_lt_two_pow (numLowBits:=ℓ) (n:=j.val)
      ⟩) * intermediateNovelBasisX 𝔽q β h_ℓ_add_R_rate ⟨ℓ, by omega⟩ ⟨x, by omega⟩

    have h_sum_right : ∑ (x: Fin (2^(ℓ - ℓ))), f_right x =
      C (a ⟨Nat.getLowBits ℓ (↑j), by exact Nat.getLowBits_lt_two_pow ℓ⟩) *
    intermediateNovelBasisX 𝔽q β h_ℓ_add_R_rate ⟨ℓ, by omega⟩ 0 := by
      have h_sum_eq := Fin.sum_congr' (b:=2^(ℓ - ℓ)) (a:=1) (f:=f_right) (by omega)
      rw [←h_sum_eq]
      rw [Fin.sum_univ_one]
      unfold f_right
      simp only [Fin.isValue, Fin.cast_zero, Fin.coe_ofNat_eq_mod, tsub_self, pow_zero,
        Nat.zero_mod, Nat.zero_shiftLeft, Nat.zero_or]
      congr

    rw [h_sum_right]

    set f_left: Fin (ℓ + R_rate - ℓ) → L := fun x =>
      if Nat.getBit (x.val) (j.val / 2 ^ ℓ) = 1 then
        eval (β ⟨ℓ + x.val, by omega⟩) (normalizedW 𝔽q β ⟨ℓ, by omega⟩)
      else 0

    simp only [eval_mul, eval_C]

    have h_eval : eval (Finset.univ.sum f_left) (intermediateNovelBasisX 𝔽q β h_ℓ_add_R_rate
      ⟨ℓ, by omega⟩ 0) = 1 := by
      have h_base_novel_basis := base_intermediateNovelBasisX 𝔽q β
        h_ℓ_add_R_rate ⟨ℓ, by exact Nat.lt_two_pow_self⟩
      simp only [intermediateNovelBasisX, Fin.coe_ofNat_eq_mod, tsub_self, pow_zero,
        Nat.zero_mod]

      set f_inner : Fin (ℓ - ℓ) → L[X] := fun x => intermediateNormVpoly 𝔽q β h_ℓ_add_R_rate
        ⟨ℓ, by omega⟩ ⟨x, by simp only; omega⟩ ^ Nat.getBit (x.val) 0

      have h_sum_eq := Fin.prod_congr' (b:=ℓ - ℓ) (a:=0) (f:=f_inner) (by omega)
      simp_rw [←h_sum_eq, Fin.prod_univ_zero]
      simp only [eval_one]

    rw [h_eval, mul_one]
    simp only [Nat.getLowBits_eq_mod_two_pow]