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

import ArkLib.Data.FieldTheory.AdditiveNTT.AdditiveNTT

import ArkLib.Data.FieldTheory.BinaryField.Tower.Impl

namespace AdditiveNTT

open ConcreteBinaryTower

section HelperFunctions

end HelperFunctions

variable {r : ℕ} [NeZero r]

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

variable {𝔽q : Type} [Field 𝔽q] [Fintype 𝔽q] [DecidableEq 𝔽q]

variable [hFq_card : Fact (Fintype.card 𝔽q = 2)]

variable [h_Fq_char_prime : Fact (Nat.Prime (ringChar 𝔽q))]

variable [Algebra 𝔽q L]

variable (β : Fin r → L) [hβ_lin_indep : Fact (LinearIndependent 𝔽q β)]

variable [h_β₀_eq_1 : Fact (β 0 = 1)]

section Algorithm

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

def bitsToU (i : Fin r) (k : Fin (2 ^ i.val)) :
    AdditiveNTT.U (L := L) (𝔽q := 𝔽q) (β := β) i :=
  let val := (Finset.univ : Finset (Fin i)).sum fun j =>
    if (Nat.getBit (n := k.val) (k := j.val) == 1) then
      β ⟨j, by admit /- proof elided -/
      ⟩
    else 0

  
  ⟨val, by admit /- proof elided -/
  ⟩

def getUElements (i : Fin r) : List L :=
  (List.finRange (2^i.val)).map fun k =>
    (Finset.univ : Finset (Fin i)).sum fun j =>
      if Nat.getBit (n := k.val) (k := j.val) == 1 then
        β ⟨j.val, by admit /- proof elided -/
        ⟩
      else 0

def evalWAt (i : Fin r) (x : L) : L :=
  ((getUElements (β := β) (ℓ := ℓ) (R_rate := R_rate) i).map (fun u => x - u)).prod

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

content: def split {k : ℕ} (h : k > 0) (x : ConcreteBTField k) :
    ConcreteBTField (k - 1) × ConcreteBTField (k - 1) :=
  let lo_bits : BitVec (2 ^ (k - 1) - 1 - 0 + 1) :=
    BitVec.extractLsb (hi := 2 ^ (k - 1) - 1) (lo := 0) x
  let hi_bits : BitVec (2 ^ k - 1 - 2 ^ (k - 1) + 1) :=
    BitVec.extractLsb (hi := 2 ^ k - 1) (lo := 2 ^ (k - 1)) x
  have h_lo : 2 ^ (k - 1) - 1 - 0 + 1 = 2 ^ (k - 1) := by admit /- proof elided -/

content: instance BitVec.instDCast : DCast Nat BitVec where
  dcast h := BitVec.cast h
  dcast_id := by admit /- proof elided -/

content: class DCast (α : Sort*) (β : α → Sort*) where
  dcast : ∀ {a a' : α}, a = a' → β a → β a'
  dcast_id : ∀ {a : α}, dcast (Eq.refl a) = id

content: instance ConcreteBTField.instDCast_local : DCast ℕ ConcreteBTField where
  dcast h_k_eq term_k1 := BitVec.cast (congrArg (fun n => 2 ^ n) h_k_eq) term_k1
  dcast_id := by admit /- proof elided -/

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

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

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

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

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

content: lemma getBit_of_lt_two_pow {n: ℕ} (a: Fin (2^n)) (k: ℕ) :
  getBit k a = if k < n then getBit k a else 0

content: lemma getBit_eq_testBit (k n : Nat) : getBit k n = if n.testBit k then 1 else 0

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

content: lemma U_card (i : Fin r) :
    Fintype.card (U 𝔽q β i) = (Fintype.card 𝔽q)^i.val

content: omit [Fintype L] [Fintype 𝔽q] h_Fq_char_prime in
lemma finrank_U (i : Fin r) :
  Module.finrank 𝔽q (U 𝔽q β i) = i

-- Local definitions used:
content: def bitsToU (i : Fin r) (k : Fin (2 ^ i.val)) :
    AdditiveNTT.U (L := L) (𝔽q := 𝔽q) (β := β) i :=
  let val := (Finset.univ : Finset (Fin i)).sum fun j =>
    if (Nat.getBit (n := k.val) (k := j.val) == 1) then
      β ⟨j, by admit /- proof elided -/
      ⟩
    else 0

  
  ⟨val, by admit /- proof elided -/
  ⟩

content: def getUElements (i : Fin r) : List L :=
  (List.finRange (2^i.val)).map fun k =>
    (Finset.univ : Finset (Fin i)).sum fun j =>
      if Nat.getBit (n := k.val) (k := j.val) == 1 then
        β ⟨j.val, by admit /- proof elided -/
        ⟩
      else 0

content: def evalWAt (i : Fin r) (x : L) : L :=
  ((getUElements (β := β) (ℓ := ℓ) (R_rate := R_rate) i).map (fun u => x - u)).prod

-- Local lemmas used:
content: lemma List.prod_finRange_eq_finset_prod {M : Type*} [CommMonoid M] {n : ℕ} (f : Fin n → M) :
    ((List.finRange n).map f).prod = ∏ i : Fin n, f i

content: theorem bitsToU_bijective (i : Fin r) :
  Function.Bijective (bitsToU (𝔽q := 𝔽q) (β := β) (ℓ := ℓ) (R_rate := R_rate) i)

-- Target theorem and proof:
theorem evalWAt_eq_W (i : Fin r) (x : L) :
  evalWAt (β := β) (ℓ := ℓ) (R_rate := R_rate) (i := i) x =
    (W (𝔽q := 𝔽q) (β := β) (i := i)).eval x
:= by
  -- 1. Convert implementation to mathematical product over Fin(2^i)
  unfold evalWAt getUElements
  rw [List.map_map]
  rw [List.prod_finRange_eq_finset_prod] -- Now the pattern matches!
  -- 2. Prepare RHS
  rw [AdditiveNTT.W, Polynomial.eval_prod]
  simp only [Polynomial.eval_sub, Polynomial.eval_X, Polynomial.eval_C]
  -- 3. Use Finset.prod_bij to show equality via the bijection
  -- LHS: ∏ k : Fin(2^i), (x - bitsToU k), RHS: ∏ u : U i,      (x - u)
  apply Finset.prod_bij (s := ((Finset.univ (α := (Fin (2^(i.val)))))))
    (t := (Finset.univ : Finset (U 𝔽q β i)))
    (i := fun k _ =>
      bitsToU (𝔽q := 𝔽q) (β := β) (ℓ := ℓ) (r := r) (R_rate := R_rate) (L := L) (i := i) k)
    (hi := by
      -- Proof that the map lands in the target set (Finset.univ)
      intro a _
      exact Finset.mem_univ _)
    (i_inj := by
      -- Proof of Injectivity (uses our previous theorem)
      intro a₁ _ a₂ _ h_eq
      exact (bitsToU_bijective (𝔽q := 𝔽q) (β := β) (ℓ := ℓ)
        (r := r) (R_rate := R_rate) (L := L) (i := i)).1 h_eq)
    (i_surj := by
      -- Proof of Surjectivity (uses our previous theorem)
      intro b _
      -- We need to provide the element 'a' and the proof it is in the set
      obtain ⟨a, ha_eq⟩ := (bitsToU_bijective (𝔽q := 𝔽q)
        (β := β) (ℓ := ℓ) (r := r) (R_rate := R_rate) (L := L) (i := i)).2 b
      use a
      constructor
      · exact ha_eq
      · exact Finset.mem_univ a
    )
    (h := by -- Proof that the values are equal: (x - bitsToU k) = (x - (bitsToU k))
      intro a ha_univ
      rfl
    )