-- 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 -/
    ⟩)

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

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

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

noncomputable def twiddleFactor (i : Fin ℓ) (u : Fin (2 ^ (ℓ + R_rate - i - 1))) : L :=
  ∑ (⟨k, hk⟩: Fin (ℓ + R_rate - i - 1)),
    if Nat.getBit k u.val = 1 then
      
        
      (normalizedW 𝔽q β ⟨i, by admit /- proof elided -/
      ⟩).eval (β ⟨i + 1 + 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) ℓ)))

noncomputable def NTTStage (i : Fin ℓ) (b : Fin (2 ^ (ℓ + R_rate)) → L) :
    Fin (2^(ℓ + R_rate)) → L :=
  have h_2_pow_i_lt_2_pow_ℓ_add_R_rate: 2^i.val < 2^(ℓ + R_rate) := by admit /- proof elided -/

noncomputable def additiveNTT (a : Fin (2 ^ ℓ) → L) : Fin (2^(ℓ + R_rate)) → L :=
  let b: Fin (2^(ℓ + R_rate)) → L := tileCoeffs a 
  Fin.foldl (n:=ℓ) (f:= fun current_b i  =>
    NTTStage 𝔽q β h_ℓ_add_R_rate (i := ⟨ℓ - 1 - i, by admit /- proof elided -/
    ⟩) current_b
  ) (init:=b)

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: notation:70 s:70 " ^^ " t:71 => Fintype.piFinset fun (i : t) ↦ s i

content: macro_rules (kind := mvEval)
  | `($p⸨$x⸩) => `(MvPolynomial.eval ($x ∘ Fin.cast (by admit /- proof elided -/
  )) $p)
  | `($p⸨$x, $y⸩) => `(MvPolynomial.eval (Fin.append $x $y ∘ Fin.cast (by admit /- proof elided -/
  )) $p)
  | `($p⸨$x, $y, $z⸩) =>
      `(MvPolynomial.eval (Fin.append (Fin.append $x $y) $z ∘ Fin.cast (by admit /- proof elided -/
      )) $p)

content: macro_rules (kind := mvEval')
  | `($p⸨$x⸩'$h) => `(MvPolynomial.eval ($x ∘ Fin.cast $h) $p)
  | `($p⸨$x, $y⸩'$h) => `(MvPolynomial.eval (Fin.append $x $y ∘ Fin.cast $h) $p)
  | `($p⸨$x, $y, $z⸩'$h) =>
      `(MvPolynomial.eval (Fin.append (Fin.append $x $y) $z ∘ Fin.cast $h) $p)

example : (X 0 + X 1 * X 2 : ℕ[X Fin 3]) ⸨![1, 2], ![8], ![]⸩ = 17 := by admit /- proof elided -/

content: macro_rules (kind := mvEvalToPolynomial)
  | `($p⸨X ⦃$i⦄, $x⸩) =>
      `(Polynomial.map (MvPolynomial.eval ($x ∘ Fin.cast (by admit /- proof elided -/
      )))
        (MvPolynomial.finSuccEquivNth _ $i $p))
  | `($p⸨X ⦃$i⦄, $x, $y⸩) =>
      `(Polynomial.map (MvPolynomial.eval (Fin.append $x $y ∘ Fin.cast (by admit /- proof elided -/
      )))
        (MvPolynomial.finSuccEquivNth _ $i $p))
  | `($p⸨X ⦃$i⦄, $x, $y, $z⸩) =>
      `(Polynomial.map (MvPolynomial.eval (Fin.append (Fin.append $x $y) $z ∘ Fin.cast (by admit /- proof elided -/
      )))
        (MvPolynomial.finSuccEquivNth _ $i $p))

content: macro_rules (kind := mvEvalToPolynomial')
  | `($p⸨X ⦃$i⦄, $x⸩'$h) =>
      `(Polynomial.map (MvPolynomial.eval ($x ∘ Fin.cast $h))
        (MvPolynomial.finSuccEquivNth _ $i $p))
  | `($p⸨X ⦃$i⦄, $x, $y⸩'$h) =>
      `(Polynomial.map (MvPolynomial.eval (Fin.append $x $y ∘ Fin.cast $h))
        (MvPolynomial.finSuccEquivNth _ $i $p))
  | `($p⸨X ⦃$i⦄, $x, $y, $z⸩'$h) =>
      `(Polynomial.map (MvPolynomial.eval (Fin.append (Fin.append $x $y) $z ∘ Fin.cast $h))
        (MvPolynomial.finSuccEquivNth _ $i $p))

example {a b n : ℕ} (x : Fin a → ℕ) (y : Fin b → ℕ) (p : ℕ[X Fin n]) (h : a + b + 1 = n) :
  p ⸨x, ![n], y⸩ =
    MvPolynomial.eval (Fin.append (Fin.append x ![n]) y ∘ Fin.cast (by admit /- proof elided -/
    )) p := rfl

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

content: def finSuccEquivNth : MvPolynomial (Fin (n + 1)) R ≃ₐ[R] Polynomial (MvPolynomial (Fin n) R) :=
  (renameEquiv R (_root_.finSuccEquiv' p)).trans (optionEquivLeft R (Fin n))

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

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: noncomputable def polynomialFromNovelCoeffs (ℓ : ℕ) (h_ℓ : ℓ ≤ r)
  (a : Fin (2 ^ ℓ) → L) : L[X] := ∑ j, C (a j) * (Xⱼ 𝔽q β ℓ h_ℓ j)

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: def getHighBits_no_shl (numLowBits : ℕ) (n : ℕ) : ℕ := n >>> numLowBits

content: def getHighBits (numLowBits : ℕ) (n : ℕ) : ℕ :=
  (getHighBits_no_shl numLowBits n) <<< numLowBits

content: def g (n : ℕ) (c : ℕ) (x : ℕ) := (x * x + c) % n

content: noncomputable def polyEvalLinearMap {L 𝔽q : Type*} [Field L] [Field 𝔽q] [Algebra 𝔽q L]
  (p : L[X]) (hp_add : IsLinearMap 𝔽q (fun x : L => p.eval x)) : L →ₗ[𝔽q] L :=
{
  toFun    := fun x => p.eval x,
  map_add' := hp_add.map_add,
  map_smul' := hp_add.map_smul
}

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: lemma getBit_eq_succ_getBit_of_mul_two_add_one {n k : ℕ} : getBit (k+1) (2*n + 1) = getBit k n

content: lemma getBit_of_shiftRight {n p : ℕ}:
  ∀ k, getBit k (n >>> p) = getBit (k+p) n

content: theorem mul_two_add_bit_lt_two_pow (a b c : ℕ) (i : Fin 2)
    (h_a : a < 2 ^ b) (h_b : b < c) :
    a * 2 + i.val < 2^c

content: lemma lt_two_pow_of_lt_two_pow_exp_le (x i j: ℕ)
    (h_x_lt_2_pow_i: x < 2^i) (h_i_le_j: i ≤ j): x < 2^j

content: lemma getBit_zero_of_two_mul {n : ℕ} : getBit 0 (2*n) = 0

content: lemma getBit_eq_succ_getBit_of_mul_two {n k : ℕ} : getBit (k+1) (2*n) = getBit k n

content: theorem Fin.sum_univ_odd_even {n : ℕ} {M : Type*} [AddCommMonoid M] (f : ℕ → M) :
    (∑ i : Fin (2 ^ n), f (2 * i)) + (∑ i : Fin (2 ^ n), f (2 * i + 1))
    = ∑ i: Fin (2 ^ (n+1)), f i

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: theorem frobenius_identity_in_algebra [Fact (Nat.Prime (ringChar Fq))]
  (f g : L[X]) : (f + g)^(Fintype.card Fq) = f^(Fintype.card Fq) + g^(Fintype.card Fq)

content: theorem linear_map_of_comp_to_linear_map_of_eval (f : L[X])
  (h_f_linear : IsLinearMap (R := Fq) (M := L[X]) (M₂ := L[X])
    (f := fun inner_p ↦ f.comp inner_p)) :
    IsLinearMap (R := Fq) (M := L) (M₂ := L) (f := fun x ↦ f.eval x)

content: lemma getBit_lt_2 {k n : Nat} : getBit k n < 2

content: theorem div_two_pow_lt_two_pow (x i j : ℕ) (h_x_lt_2_pow_i : x < 2 ^ (i + j)): x / 2^j < 2^(i)

content: lemma sum_of_and_eq_zero_is_or {n m : ℕ} (h_n_AND_m : n &&& m = 0) : n + m = n ||

content: lemma sum_eq_xor_plus_twice_and (n : Nat) : ∀ m : ℕ, n + m = (n ^^^ m) + 2 * (n &&& m)

content: lemma and_by_split_lowBits {n m n1 m1 bn bm : ℕ} (h_bn : bn < 2) (h_bm : bm < 2)
  (h_n : n = n1 * 2 + bn) (h_m : m = m1 * 2 + bm):
  n &&& m = (n1 &&& m1) * 2 + (bn &&& bm)

content: lemma div_2_form {nD2 b : ℕ} (h_b : b < 2):
  (nD2 * 2 + b) / 2 = nD2

content: lemma or_by_split_lowBits {n m n1 m1 bn bm : ℕ} (h_bn : bn < 2) (h_bm : bm < 2)
  (h_n : n = n1 * 2 + bn) (h_m : m = m1 * 2 + bm):
  n ||

content: lemma xor_by_split_lowBits {n m n1 m1 bn bm : ℕ} (h_bn : bn < 2) (h_bm : bm < 2)
  (h_n : n = n1 * 2 + bn) (h_m : m = m1 * 2 + bm):
  n ^^^ m = (n1 ^^^ m1) * 2 + (bn ^^^ bm)

content: lemma xor_of_and_eq_zero_is_or {n m : ℕ} (h_n_AND_m : n &&& m = 0) : n ^^^ m = n ||

content: lemma eq_iff_eq_all_getBits {n m : ℕ} : n = m ↔ ∀ k, getBit k n = getBit k m

content: lemma and_eq_zero_iff_and_each_getBit_eq_zero {n m : ℕ} :
    n &&& m = 0 ↔ ∀ k, Nat.getBit k n &&& Nat.getBit k m = 0

content: lemma shiftRight_and_one_distrib {n m k : ℕ} :
    Nat.getBit k (n &&& m) = Nat.getBit k n &&& Nat.getBit k m

content: lemma and_two_pow_eq_zero_of_getBit_0 {n i : ℕ} (h_getBit : getBit i n = 0)
    : n &&& (2 ^ i) = 0

content: lemma getBit_two_pow {i k : ℕ} : (getBit k (2^i) = if i == k then 1 else 0)

content: lemma getLowBits_succ {n: ℕ} (numLowBits: ℕ) :
    getLowBits (numLowBits + 1) n = getLowBits numLowBits n
    + (getBit numLowBits n) <<< numLowBits

content: lemma getBit_of_lowBits {n: ℕ} (numLowBits : ℕ) : ∀ k, getBit k (getLowBits numLowBits n) =
    if k < numLowBits then getBit k n else 0

content: lemma getBit_of_and {n m k: ℕ} : getBit k (n &&& m) = getBit k n &&& getBit k m

content: lemma getBit_of_two_pow_sub_one {i k: ℕ} : getBit k (2^i - 1) =
    if k < i then 1 else 0

content: lemma getBit_of_or {n m k: ℕ} : getBit k (n ||| m) = getBit k n ||

content: lemma sum_of_and_eq_zero_is_xor {n m : ℕ} (h_n_AND_m : n &&& m = 0) : n + m = n ^^^ m

content: lemma xor_eq_sub_iff_submask {n m : ℕ} (h: m ≤ n) : n ^^^ m = n - m ↔ n &&& m = m

content: lemma getBit_of_xor {n m k: ℕ} : getBit k (n ^^^ m) = getBit k n ^^^ getBit k m

content: lemma normalizedWᵢ_eval_βᵢ_eq_1 {i : Fin r} :
    (normalizedW (𝔽q := 𝔽q) (β := β) (i :=i)).eval (β i) = 1

content: lemma num_eq_highBits_xor_lowBits {n: ℕ} (numLowBits: ℕ) :
  n = getHighBits numLowBits n ^^^ getLowBits numLowBits n

content: theorem and_highBits_lowBits_eq_zero {n : ℕ} (numLowBits : ℕ) :
    getHighBits numLowBits n &&& getLowBits numLowBits n = 0

content: lemma getBit_of_shiftLeft {n p : ℕ}:
  ∀ k, getBit (k) (n <<< p) = if k < p then 0 else getBit (k - p) n

content: lemma getBit_of_multiple_of_power_of_two {n p : ℕ}: ∀ k,
  getBit (k) (2^p * n) = if k < p then 0 else getBit (k-p) n

content: lemma num_eq_highBits_add_lowBits {n: ℕ} (numLowBits: ℕ) :
  n = getHighBits numLowBits n + getLowBits numLowBits n

content: lemma getBit_of_add_distrib {n m k : ℕ}
  (h_n_AND_m : n &&& m = 0) : getBit k (n + m) = getBit k n + getBit k m

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

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

content: lemma and_two_pow_eq_two_pow_of_getBit_1 {n i : ℕ} (h_getBit: getBit i n = 1) :
    n &&& (2 ^ i) = 2 ^ i

content: lemma add_two_pow_of_getBit_eq_zero_lt_two_pow {n m i : ℕ} (h_n: n < 2^m) (h_i: i < m)
  (h_getBit_at_i_eq_zero: getBit i n = 0) :
  n + 2^i < 2^m

content: lemma getLowBits_le_self {n : ℕ} (numLowBits : ℕ) : getLowBits numLowBits n ≤ n

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_zero_eq_zero {n : ℕ} : getLowBits 0 n = 0

-- 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 evenRefinement (i : Fin (ℓ))
    (coeffs : Fin (2 ^ (ℓ - i)) → L) : L[X] :=
  ∑ (⟨j, hj⟩: Fin (2^(ℓ-i-1))), C (coeffs ⟨j*2, by admit /- proof elided -/
  ⟩) * (intermediateNovelBasisX 𝔽q β h_ℓ_add_R_rate ⟨i+1, by admit /- proof elided -/
  ⟩ ⟨j, hj⟩)

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

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: noncomputable def twiddleFactor (i : Fin ℓ) (u : Fin (2 ^ (ℓ + R_rate - i - 1))) : L :=
  ∑ (⟨k, hk⟩: Fin (ℓ + R_rate - i - 1)),
    if Nat.getBit k u.val = 1 then
      
        
      (normalizedW 𝔽q β ⟨i, by admit /- proof elided -/
      ⟩).eval (β ⟨i + 1 + 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: noncomputable def NTTStage (i : Fin ℓ) (b : Fin (2 ^ (ℓ + R_rate)) → L) :
    Fin (2^(ℓ + R_rate)) → L :=
  have h_2_pow_i_lt_2_pow_ℓ_add_R_rate: 2^i.val < 2^(ℓ + R_rate) := by admit /- proof elided -/

content: noncomputable def additiveNTT (a : Fin (2 ^ ℓ) → L) : Fin (2^(ℓ + R_rate)) → L :=
  let b: Fin (2^(ℓ + R_rate)) → L := tileCoeffs a 
  Fin.foldl (n:=ℓ) (f:= fun current_b i  =>
    NTTStage 𝔽q β h_ℓ_add_R_rate (i := ⟨ℓ - 1 - i, by admit /- proof elided -/
    ⟩) current_b
  ) (init:=b)

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: omit [DecidableEq L] [DecidableEq 𝔽q] h_Fq_char_prime hF₂ hβ_lin_indep h_β₀_eq_1 in
theorem qMap_eval_𝔽q_eq_0 (i : Fin r) :
  ∀ c: 𝔽q, (qMap 𝔽q β i).eval (algebraMap 𝔽q L c) = 0

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: theorem qMap_is_linear_map (i : Fin r) :
  IsLinearMap 𝔽q (f:=fun inner_p ↦ (qMap 𝔽q β i).comp inner_p)

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: omit [Fintype L] [DecidableEq L] in
theorem Polynomial.foldl_comp (n : ℕ) (f : Fin n → L[X]) : ∀ initInner initOuter: L[X],
    Fin.foldl (n:=n) (fun acc j => (f j).comp acc) (initOuter.comp initInner)
    = (Fin.foldl (n:=n) (fun acc j => (f j).comp acc) (initOuter)).comp initInner

content: theorem intermediateNormVpoly_comp_qmap (i : Fin (ℓ))
    (k : Fin (ℓ - i - 1)) : -- corresponds to intermediateNormVpoly_comp

content: theorem intermediateNormVpoly_comp_qmap_helper (i : Fin (ℓ))
    (k : Fin (ℓ - (↑i + 1))) :
    (intermediateNormVpoly 𝔽q β h_ℓ_add_R_rate
      ⟨↑i + 1, by omega⟩ (k:=⟨k, by simp only; omega⟩)).comp (qMap 𝔽q β ⟨↑i, by omega⟩) =
    intermediateNormVpoly 𝔽q β h_ℓ_add_R_rate
      ⟨↑i, by omega⟩ ⟨k + 1, by simp only; 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

content: lemma even_index_intermediate_novel_basis_decomposition (i : Fin ℓ) (j : Fin (2 ^ (ℓ - i - 1))) :
  intermediateNovelBasisX 𝔽q β h_ℓ_add_R_rate ⟨i, by omega⟩ ⟨j * 2, by
    apply mul_two_add_bit_lt_two_pow j (ℓ-i-1) (ℓ-i) ⟨0, by omega⟩ (by omega) (by omega)
  ⟩  = (intermediateNovelBasisX 𝔽q β h_ℓ_add_R_rate ⟨i+1, by omega⟩ ⟨j, by
    apply lt_two_pow_of_lt_two_pow_exp_le j (ℓ-i-1) (ℓ-(i+1)) (by omega) (by omega)
  ⟩).comp (qMap 𝔽q β ⟨i, by omega⟩)

content: lemma odd_index_intermediate_novel_basis_decomposition
    (i : Fin ℓ) (j : Fin (2 ^ (ℓ - i - 1))) :
    intermediateNovelBasisX 𝔽q β h_ℓ_add_R_rate ⟨i, by omega⟩ ⟨j * 2 + 1, by
      apply mul_two_add_bit_lt_two_pow j (ℓ-i-1) (ℓ-i) ⟨1, by omega⟩ (by omega) (by omega)
    ⟩  = X * (intermediateNovelBasisX 𝔽q β h_ℓ_add_R_rate ⟨i+1, by omega⟩ ⟨j, by
      apply lt_two_pow_of_lt_two_pow_exp_le j (ℓ-i-1) (ℓ-(i+1)) (by omega) (by omega)
    ⟩).comp (qMap 𝔽q β ⟨i, by omega⟩)

content: theorem evaluation_poly_split_identity (i : Fin (ℓ))
    (coeffs : Fin (2 ^ (ℓ - i)) → L) :
  let P_i: L[X] := intermediateEvaluationPoly 𝔽q β h_ℓ_add_R_rate ⟨i, by omega⟩ coeffs
  let P_even_i_plus_1: L[X] := evenRefinement 𝔽q β h_ℓ_add_R_rate i coeffs
  let P_odd_i_plus_1: L[X] := oddRefinement 𝔽q β h_ℓ_add_R_rate i coeffs
  let q_i: L[X] := qMap 𝔽q β ⟨i, by omega⟩
  P_i = (P_even_i_plus_1.comp q_i) + X * (P_odd_i_plus_1.comp q_i)

content: lemma intermediate_poly_P_base (h_ℓ : ℓ ≤ r) (coeffs : Fin (2 ^ ℓ) → L) :
  intermediateEvaluationPoly 𝔽q β h_ℓ_add_R_rate ⟨0, by omega⟩ coeffs =
    polynomialFromNovelCoeffs 𝔽q β ℓ h_ℓ coeffs

content: omit [DecidableEq L] [DecidableEq 𝔽q] [Fintype 𝔽q] h_Fq_char_prime hF₂ hβ_lin_indep h_β₀_eq_1 in
lemma evaluationPointω_eq_twiddleFactor_of_div_2 (i : Fin ℓ) (x : Fin (2 ^ (ℓ + R_rate - i))) :
  evaluationPointω 𝔽q β h_ℓ_add_R_rate ⟨i, by omega⟩ x =
  twiddleFactor 𝔽q β h_ℓ_add_R_rate ⟨i, by omega⟩ ⟨x/2, by
    have h := div_two_pow_lt_two_pow (x:=x) (i := ℓ + R_rate - i - 1) (j:=1) (by
      rw [Nat.sub_add_cancel (by omega)]; omega)
    simp only [pow_one] at h
    calc _ < 2 ^ (ℓ + R_rate - i - 1) := by omega
      _ = _

content: omit [DecidableEq 𝔽q] hF₂ h_β₀_eq_1 in
lemma eval_point_ω_eq_next_twiddleFactor_comp_qmap

  (i : Fin ℓ) (x : Fin (2 ^ (ℓ + R_rate - (i + 1)))) :
  -- `j = u||b||v` => x here means u at level i
  evaluationPointω 𝔽q β h_ℓ_add_R_rate ⟨i.val+1, by omega⟩ x =
  eval (twiddleFactor 𝔽q β h_ℓ_add_R_rate ⟨i, by omega⟩ ⟨x.val, by
    calc x.val < 2 ^ (ℓ + R_rate - (i.val + 1)) := by omega
      _ = 2 ^ (ℓ + R_rate - i.val - 1) := by rfl
  ⟩) (qMap 𝔽q β ⟨i, by omega⟩)

content: omit [NeZero r] [Field L] [Fintype L] [DecidableEq L] [DecidableEq 𝔽q] [Field 𝔽q] [Algebra 𝔽q L] in
lemma base_coeffsBySuffix (a : Fin (2 ^ ℓ) → L) :
  coeffsBySuffix (r:=r) (R_rate := R_rate) a 0 0 = a

content: theorem evenRefinement_eq_novel_poly_of_0_leading_suffix (i : Fin ℓ) (v : Fin (2 ^ i.val))
    (original_coeffs : Fin (2 ^ ℓ) → L) :
    have h_v: v.val < 2 ^ (i.val + 1) := by
      calc v.val < 2 ^ i.val := by omega
        _ < 2 ^ (i.val + 1) := by apply Nat.pow_lt_pow_right (by omega) (by omega)
    evenRefinement 𝔽q β h_ℓ_add_R_rate i (coeffsBySuffix (r:=r)
      (R_rate:=R_rate) (a:=original_coeffs) ⟨i, by omega⟩ v) =
    intermediateEvaluationPoly 𝔽q β h_ℓ_add_R_rate ⟨i + 1, by omega⟩
      (coeffsBySuffix (r:=r) (R_rate:=R_rate) original_coeffs ⟨i + 1, by omega⟩ ⟨v, h_v⟩)

content: theorem oddRefinement_eq_novel_poly_of_1_leading_suffix (i : Fin ℓ) (v : Fin (2 ^ i.val))
    (original_coeffs : Fin (2 ^ ℓ) → L) :
    have h_v: v.val ||| (1 <<< i.val) < 2 ^ (i.val + 1)

content: 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⟩)

content: lemma NTTStage_correctness (i : Fin (ℓ))
    (input_buffer : Fin (2 ^ (ℓ + R_rate)) → L) (original_coeffs : Fin (2 ^ ℓ) → L) :
    additiveNTTInvariant 𝔽q β h_ℓ_add_R_rate (evaluation_buffer:=input_buffer)
      (original_coeffs:=original_coeffs) (i := ⟨i.val+1, by omega⟩) →
    additiveNTTInvariant 𝔽q β h_ℓ_add_R_rate (evaluation_buffer:=NTTStage 𝔽q β h_ℓ_add_R_rate
      ⟨i, by omega⟩ input_buffer) (original_coeffs:=original_coeffs) ⟨i, by omega⟩

content: lemma foldl_NTTStage_inductive_aux (h_ℓ : ℓ ≤ r) (k : Fin (ℓ + 1))
    (original_coeffs : Fin (2 ^ ℓ) → L) :
    additiveNTTInvariant 𝔽q β h_ℓ_add_R_rate
    (Fin.foldl k (fun current_b i ↦ NTTStage 𝔽q β h_ℓ_add_R_rate
      ⟨ℓ - i -1, by omega⟩ current_b) (tileCoeffs original_coeffs))
    original_coeffs ⟨ℓ - k, by omega⟩

-- Target theorem and proof:
theorem additiveNTT_correctness (h_ℓ : ℓ ≤ r)
    (original_coeffs : Fin (2 ^ ℓ) → L)
    (output_buffer : Fin (2 ^ (ℓ + R_rate)) → L)
    (h_alg : output_buffer = additiveNTT 𝔽q β h_ℓ_add_R_rate original_coeffs) :
    let P := polynomialFromNovelCoeffs 𝔽q β ℓ h_ℓ original_coeffs
    ∀ (j : Fin (2^(ℓ + R_rate))),
      output_buffer j = P.eval (evaluationPointω 𝔽q β h_ℓ_add_R_rate ⟨0, by omega⟩ j)
:= polynomialFromNovelCoeffs 𝔽q β ℓ h_ℓ original_coeffs
    ∀ (j : Fin (2^(ℓ + R_rate))),
      output_buffer j = P.eval (evaluationPointω 𝔽q β h_ℓ_add_R_rate ⟨0, by omega⟩ j) :=
  by
  simp only [Fin.zero_eta]
  intro j
  simp only [h_alg]
  unfold additiveNTT
  set output_foldl := Fin.foldl ℓ (fun current_b i ↦ NTTStage 𝔽q β h_ℓ_add_R_rate
    ⟨ℓ - i -1, by omega⟩ current_b) (tileCoeffs original_coeffs)

  have output_foldl_correctness : additiveNTTInvariant 𝔽q β h_ℓ_add_R_rate
    output_foldl original_coeffs ⟨0, by omega⟩ := by
    have res := foldl_NTTStage_inductive_aux 𝔽q β h_ℓ_add_R_rate
      h_ℓ
      (k:=⟨ℓ, by omega⟩) original_coeffs
    simp only [tsub_self, Fin.zero_eta] at res
    exact res

  have h_nat_point_ω_eq_j: j.val / 2 * 2 + j.val % 2 = j := by
    have h_j_mod_2_eq_0: j.val % 2 < 2 := by omega
    exact Nat.div_add_mod' (↑j) 2

  simp only [additiveNTTInvariant] at output_foldl_correctness
  have res := output_foldl_correctness j
  unfold output_foldl at res
  simp only [Fin.zero_eta, Nat.sub_zero, pow_zero, Nat.div_one, Fin.eta,
    Nat.pow_zero, Nat.getLowBits_zero_eq_zero (n := j.val), Fin.isValue, base_coeffsBySuffix] at res
  simp only [←
    intermediate_poly_P_base 𝔽q β h_ℓ_add_R_rate
      h_ℓ original_coeffs,
    Fin.zero_eta]
  rw [←res]
  simp_rw [Nat.sub_right_comm] -- ℓ - 1 - ↑i = ℓ - ↑i - 1