-- 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 sDomain (i : Fin r) : Subspace 𝔽q L :=
  let W_i_norm := normalizedW 𝔽q β i
  let h_W_i_norm_is_additive : IsLinearMap 𝔽q (fun x : L => W_i_norm.eval x) :=
    AdditiveNTT.normalizedW_is_additive 𝔽q β i
  Submodule.map (polyEvalLinearMap W_i_norm h_W_i_norm_is_additive)
    (U 𝔽q β ⟨ℓ + R_rate, h_ℓ_add_R_rate⟩)

def sBasis (i : Fin r) (h_i : i < ℓ + R_rate) : Fin (ℓ + R_rate - i) → L :=
  fun k => β ⟨i + k.val, by admit /- proof elided -/
  ⟩

noncomputable def sDomain_basis (i : Fin r) (h_i : i < ℓ + R_rate) :
    Basis (Fin (ℓ + R_rate - i)) 𝔽q (
      sDomain 𝔽q β h_ℓ_add_R_rate i) :=

noncomputable section DomainBijection

def splitPointIntoCoeffs (i : Fin r) (h_i : i < ℓ + R_rate)
  (x : sDomain 𝔽q β h_ℓ_add_R_rate i) :
  Fin (ℓ + R_rate - i.val) → ℕ := fun j =>
    if ((sDomain_basis 𝔽q β
    h_ℓ_add_R_rate i h_i).repr x j = 0) then
      0 else 1

noncomputable def sDomainToFin (i : Fin r) (h_i : i < ℓ + R_rate)
  (x : sDomain 𝔽q β h_ℓ_add_R_rate i) :
  Fin (2^(ℓ + R_rate - i.val)) :=

def finToBinaryCoeffs (i : Fin r) (idx : Fin (2 ^ (ℓ + R_rate - i.val))) :
  Fin (ℓ + R_rate - i.val) → 𝔽q := fun j =>
    if (Nat.getBit (k:=j) (n:=idx)) = 1 then (1 : 𝔽q) else (0 : 𝔽q)

-- Repository definitions used:
content: def getBit (k n : Nat) : Nat := (n >>> k) &&& 1

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

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 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: def binaryFinMapToNat {n : ℕ} (m : Fin n → ℕ) (h_binary : ∀ j: Fin n, m j ≤ 1) : Fin (2^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_of_binaryFinMapToNat {n : ℕ} (m : Fin n → ℕ) (h_binary: ∀ j: Fin n, m j ≤ 1) :
    ∀ k: ℕ, Nat.getBit k (binaryFinMapToNat m h_binary).val
      = if h_k: k < n then m ⟨k, by omega⟩ else 0

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

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 eq_iff_eq_all_getBits {n m : ℕ} : n = m ↔ ∀ k, getBit k n = getBit k m

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

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

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 getBit_of_xor {n m k: ℕ} : getBit k (n ^^^ m) = getBit k n ^^^ getBit k m

content: lemma getBit_zero_eq_zero {k : Nat} : getBit k 0 = 0

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

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)

-- Local definitions used:
content: noncomputable def sDomain (i : Fin r) : Subspace 𝔽q L :=
  let W_i_norm := normalizedW 𝔽q β i
  let h_W_i_norm_is_additive : IsLinearMap 𝔽q (fun x : L => W_i_norm.eval x) :=
    AdditiveNTT.normalizedW_is_additive 𝔽q β i
  Submodule.map (polyEvalLinearMap W_i_norm h_W_i_norm_is_additive)
    (U 𝔽q β ⟨ℓ + R_rate, h_ℓ_add_R_rate⟩)

content: def sBasis (i : Fin r) (h_i : i < ℓ + R_rate) : Fin (ℓ + R_rate - i) → L :=
  fun k => β ⟨i + k.val, by admit /- proof elided -/
  ⟩

content: noncomputable def sDomain_basis (i : Fin r) (h_i : i < ℓ + R_rate) :
    Basis (Fin (ℓ + R_rate - i)) 𝔽q (
      sDomain 𝔽q β h_ℓ_add_R_rate i) :=

content: def splitPointIntoCoeffs (i : Fin r) (h_i : i < ℓ + R_rate)
  (x : sDomain 𝔽q β h_ℓ_add_R_rate i) :
  Fin (ℓ + R_rate - i.val) → ℕ := fun j =>
    if ((sDomain_basis 𝔽q β
    h_ℓ_add_R_rate i h_i).repr x j = 0) then
      0 else 1

content: noncomputable def sDomainToFin (i : Fin r) (h_i : i < ℓ + R_rate)
  (x : sDomain 𝔽q β h_ℓ_add_R_rate i) :
  Fin (2^(ℓ + R_rate - i.val)) :=

content: def finToBinaryCoeffs (i : Fin r) (idx : Fin (2 ^ (ℓ + R_rate - i.val))) :
  Fin (ℓ + R_rate - i.val) → 𝔽q := fun j =>
    if (Nat.getBit (k:=j) (n:=idx)) = 1 then (1 : 𝔽q) else (0 : 𝔽q)

-- Local lemmas used:
content: omit h_Fq_char_prime in
lemma 𝔽q_element_eq_zero_or_eq_one : ∀ c: 𝔽q, c = 0 ∨ c = 1

-- Target theorem and proof:
omit h_β₀_eq_1 in
lemma finToBinaryCoeffs_sDomainToFin (i : Fin r) (h_i : i < ℓ + R_rate)
    (x : sDomain 𝔽q β h_ℓ_add_R_rate i) :
    let pointFinIdx := (sDomainToFin 𝔽q β h_ℓ_add_R_rate i h_i) x
    finToBinaryCoeffs 𝔽q (i := i) (idx :=pointFinIdx) =
    (sDomain_basis 𝔽q β
    h_ℓ_add_R_rate i h_i).repr x
:= (sDomainToFin 𝔽q β h_ℓ_add_R_rate i h_i) x
    finToBinaryCoeffs 𝔽q (i := i) (idx :=pointFinIdx) =
    (sDomain_basis 𝔽q β
    h_ℓ_add_R_rate i h_i).repr x:= by
  simp only
  ext j
  -- Unfold the definitions to get to the core logic
  dsimp [sDomainToFin, finToBinaryCoeffs, splitPointIntoCoeffs]
  -- `Nat.getBit` is the inverse of `Nat.binaryFinMapToNat`
  rw [Nat.getBit_of_binaryFinMapToNat]
  -- Let `c` be the j-th coefficient we are considering
  set c := (sDomain_basis 𝔽q β
    h_ℓ_add_R_rate i h_i).repr x j
  -- Since the field has card 2, `c` must be 0 or 1
  have hc : c = 0 ∨ c = 1 := by exact 𝔽q_element_eq_zero_or_eq_one 𝔽q c
    -- exact ((Fintype.card_eq_two_iff _).mp h_Fq_card_eq_2).right c
  -- We can now split on whether c is 0 or 1
  rcases hc with h_c_zero | h_c_one
  · -- Case 1: The coefficient is 0
    simp only [Fin.is_lt, ↓reduceDIte, Fin.eta, h_c_zero, ite_eq_right_iff, one_ne_zero, imp_false,
      ne_eq]
    unfold splitPointIntoCoeffs
    simp only [ite_eq_right_iff, zero_ne_one, imp_false, Decidable.not_not]
    omega
  · -- Case 2: The coefficient is 1
    simp only [Fin.is_lt, ↓reduceDIte, Fin.eta, h_c_one, ite_eq_left_iff, zero_ne_one, imp_false,
      Decidable.not_not]
    unfold splitPointIntoCoeffs
    simp only [ite_eq_right_iff, zero_ne_one, imp_false, ne_eq]
    change ¬(c) = 0
    rw [h_c_one]
    exact one_ne_zero