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

-- 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 g (n : ℕ) (c : ℕ) (x : ℕ) := (x * x + c) % n

-- 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

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

-- Local lemmas used:
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: 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⟩)

-- Target theorem and proof:
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)
:= 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) := by

  simp only [intermediateEvaluationPoly, Fin.eta]
  simp only [evenRefinement, Fin.eta, sum_comp, mul_comp, C_comp, oddRefinement]

  set leftEvenTerm := ∑ ⟨j, hj⟩ : Fin (2 ^ (ℓ - ↑i - 1)), C (coeffs ⟨j * 2, by
    exact 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, by omega⟩ ⟨j * 2, by
    exact mul_two_add_bit_lt_two_pow j (ℓ-i-1) (ℓ-i) ⟨0, by omega⟩ (by omega) (by omega)
  ⟩
  set leftOddTerm := ∑ ⟨j, hj⟩ : Fin (2 ^ (ℓ - ↑i - 1)), C (coeffs ⟨j * 2 + 1, by
    apply mul_two_add_bit_lt_two_pow j (ℓ-i-1) (ℓ-i) ⟨1, by omega⟩ (by omega) (by omega)
  ⟩) * intermediateNovelBasisX 𝔽q β h_ℓ_add_R_rate ⟨↑i, by omega⟩ ⟨j * 2 + 1, by
    exact mul_two_add_bit_lt_two_pow j (ℓ-i-1) (ℓ-i) ⟨1, by omega⟩ (by omega) (by omega)
  ⟩

  have h_split_P_i: ∑ ⟨j, hj⟩ : Fin (2 ^ (ℓ - ↑i)), C (coeffs ⟨j, by
    apply lt_two_pow_of_lt_two_pow_exp_le j (ℓ-i) (ℓ-i) (by omega) (by omega)
  ⟩) * intermediateNovelBasisX 𝔽q β h_ℓ_add_R_rate ⟨↑i, by omega⟩ ⟨j, by omega⟩ =
  leftEvenTerm + leftOddTerm
  := by
    unfold leftEvenTerm leftOddTerm
    simp only [Fin.eta]

    -- ⊢ ∑ k ∈ Fin (2 ^ (ℓ - ↑i)), C (coeffsₖ) * Xₖ⁽ⁱ⁾(X) = -- just pure even odd split
    -- ∑ k ∈ Fin (2 ^ (ℓ - ↑i - 1)), C (coeffs₂ₖ) * X₂ₖ⁽ⁱ⁾(X) +
    -- ∑ k ∈ Fin (2 ^ (ℓ - ↑i - 1)), C (coeffs₂ₖ+1) * X₂ₖ+1⁽ⁱ⁾(X)

    set f1 := fun x: ℕ => -- => use a single function to represent the sum
      if hx: x < 2 ^ (ℓ - ↑i) then
        C (coeffs ⟨x, hx⟩) *
          intermediateNovelBasisX 𝔽q β h_ℓ_add_R_rate ⟨↑i, by omega⟩ ⟨x, by omega⟩
      else 0

    have h_x: ∀ x: Fin (2 ^ (ℓ - ↑i)), f1 x.val =
      C (coeffs ⟨x.val, by omega⟩) *
        intermediateNovelBasisX 𝔽q β h_ℓ_add_R_rate ⟨↑i, by omega⟩
          ⟨x.val, by simp only; omega⟩ := by
      intro x
      unfold f1
      simp only [Fin.is_lt, ↓reduceDIte, Fin.eta]

    conv_lhs =>
      enter [2, x]
      rw [←h_x x]

    have h_x_2: ∀ x: Fin (2 ^ (ℓ - ↑i - 1)), f1 (x*2) =
      C (coeffs ⟨x.val * 2, by
        calc _ < 2 ^ (ℓ - i - 1) * 2 := by omega
          _ = 2 ^ (ℓ - i) := Nat.two_pow_pred_mul_two (w:=ℓ - i) (h:=by omega)
      ⟩) *
        intermediateNovelBasisX 𝔽q β h_ℓ_add_R_rate ⟨↑i, by omega⟩ ⟨x.val * 2, by
          exact mul_two_add_bit_lt_two_pow x.val (ℓ-i-1) (ℓ-i) ⟨0, by omega⟩ (by omega) (by omega)
        ⟩ := by
      intro x
      unfold f1
      simp only
      have h_x_lt_2_pow_i_minus_1 :=
        mul_two_add_bit_lt_two_pow x.val (ℓ-i-1) (ℓ-i) ⟨0, by omega⟩ (by omega) (by omega)
      simp at h_x_lt_2_pow_i_minus_1
      simp only [h_x_lt_2_pow_i_minus_1, ↓reduceDIte]

    conv_rhs =>
      enter [1, 2, x]
      rw [←h_x_2 x]

    have h_x_3: ∀ x: Fin (2 ^ (ℓ - ↑i - 1)), f1 (x*2+1) =
      C (coeffs ⟨x.val * 2 + 1, by
        calc _ < 2 ^ (ℓ - i - 1) * 2 := by omega
          _ = 2 ^ (ℓ - i) := Nat.two_pow_pred_mul_two (w:=ℓ - i) (h:=by omega)
      ⟩) *
        intermediateNovelBasisX 𝔽q β h_ℓ_add_R_rate ⟨↑i, by omega⟩ ⟨x.val * 2 + 1, by
          exact mul_two_add_bit_lt_two_pow x.val (ℓ-i-1) (ℓ-i) ⟨1, by omega⟩ (by omega) (by omega)
        ⟩ := by
      intro x
      unfold f1
      simp only
      have h_x_lt_2_pow_i_minus_1 := mul_two_add_bit_lt_two_pow x.val
        (ℓ-i-1) (ℓ-i) ⟨1, by omega⟩ (by omega) (by omega)
      simp only [h_x_lt_2_pow_i_minus_1, ↓reduceDIte]

    conv_rhs =>
      enter [2, 2, x]
      rw [←h_x_3 x]

    -- ⊢ ∑ x, f1 ↑x = ∑ x, f1 (↑x * 2) + ∑ x, f1 (↑x * 2 + 1)

    have h_1: ∑ i ∈ Finset.range (2 ^ (ℓ - ↑i)), f1 i
      = ∑ i ∈ Finset.range (2 ^ (ℓ - ↑i - 1 + 1)), f1 i := by
      congr
      omega

    have res := Fin.sum_univ_odd_even (f:=f1) (n:=(ℓ - ↑i - 1))
    conv_rhs at res =>
      rw [Fin.sum_univ_eq_sum_range]
      rw [←h_1]
      rw [←Fin.sum_univ_eq_sum_range]

    rw [←res]
    congr
    · funext i
      rw [mul_comm]
    · funext i
      rw [mul_comm]

  conv_lhs => rw [h_split_P_i]

  set rightEvenTerm := ∑ ⟨j, hj⟩ : Fin (2 ^ (ℓ - ↑i - 1)),
      C (coeffs ⟨j * 2, by
        calc _ < 2 ^ (ℓ - i - 1) * 2 := by omega
          _ = 2 ^ (ℓ - i) := Nat.two_pow_pred_mul_two (w:=ℓ - i) (h:=by omega)
      ⟩) *
        (intermediateNovelBasisX 𝔽q β h_ℓ_add_R_rate ⟨i + 1, by omega⟩ ⟨j, by
          apply lt_two_pow_of_lt_two_pow_exp_le (x:=j)
            (i := ℓ-↑i-1) (j:=ℓ-↑i-1) (by omega) (by omega)
        ⟩).comp (qMap 𝔽q β ⟨i, by omega⟩)

  set rightOddTerm :=
    X *
      ∑ ⟨j, hj⟩ : Fin (2 ^ (ℓ - ↑i - 1)),
        C (coeffs ⟨j * 2 + 1, by
          calc _ < 2 ^ (ℓ - i - 1) * 2 := by omega
            _ = 2 ^ (ℓ - i) := Nat.two_pow_pred_mul_two (w:=ℓ - i) (h:=by omega)
        ⟩) *
          (intermediateNovelBasisX 𝔽q β h_ℓ_add_R_rate ⟨i + 1, by omega⟩ ⟨j, by
            apply lt_two_pow_of_lt_two_pow_exp_le (x:=j)
              (i := ℓ-↑i-1) (j:=ℓ-↑i-1) (by omega) (by omega)
          ⟩).comp (qMap 𝔽q β ⟨i, by omega⟩)

  conv_rhs => change rightEvenTerm + rightOddTerm

  have h_right_even_term: leftEvenTerm = rightEvenTerm := by
    unfold rightEvenTerm leftEvenTerm
    apply Finset.sum_congr rfl
    intro j hj
    simp only [Fin.eta, mul_eq_mul_left_iff, map_eq_zero]
    --  X₂ⱼ⁽ⁱ⁾ = Xⱼ⁽ⁱ⁺¹⁾(q⁽ⁱ⁾(X)) ∨ a₂ⱼ = 0
    by_cases h_a_j_eq_0: coeffs ⟨j * 2, by
      calc _ < 2 ^ (ℓ - i - 1) * 2 := by omega
        _ = 2 ^ (ℓ - i) := Nat.two_pow_pred_mul_two (w:=ℓ - i) (h:=by omega)
    ⟩ = 0
    · simp only [h_a_j_eq_0, or_true]
    · simp only [h_a_j_eq_0, or_false]
      --  X₂ⱼ⁽ⁱ⁾ = Xⱼ⁽ⁱ⁺¹⁾(q⁽ⁱ⁾(X))

      exact even_index_intermediate_novel_basis_decomposition
        𝔽q β h_ℓ_add_R_rate (i := ⟨i, by omega⟩) j

  have h_right_odd_term: rightOddTerm = leftOddTerm := by
    unfold rightOddTerm leftOddTerm
    simp only [Fin.eta]
    conv_rhs =>
      simp only [Fin.is_lt, odd_index_intermediate_novel_basis_decomposition, Fin.eta]
      enter [2, x];
      rw [mul_comm (a:=X)]

    rw [Finset.mul_sum]
    congr
    funext x
    ring_nf -- just associativity and commutativity of multiplication in L[X]

  rw [h_right_even_term, h_right_odd_term]