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

-- Repository definitions used:
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: def getBit (k n : Nat) : Nat := (n >>> k) &&& 1

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

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: lemma normalizedWᵢ_vanishing (i : Fin r) :
  ∀ u ∈ U 𝔽q β i, (normalizedW 𝔽q β i).eval u = 0

content: lemma Wᵢ_vanishing (i : Fin r) :
  ∀ u ∈ U 𝔽q β i, (W 𝔽q β i).eval u = 0

content: omit hF₂ in
theorem normalizedW_is_additive (i : Fin r) :
  IsLinearMap 𝔽q (f := fun x ↦ (normalizedW 𝔽q β i).eval x)

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

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

-- Local lemmas used:
content: omit [NeZero r] [Field L] [Fintype L] [DecidableEq L] [Field 𝔽q] [Algebra 𝔽q L] in
lemma sBasis_range_eq (i : Fin r) (h_i : i < ℓ + R_rate) :
    β '' Set.Ico i ⟨ℓ + R_rate, h_ℓ_add_R_rate⟩
    = Set.range (sBasis β h_ℓ_add_R_rate i h_i)

-- Target theorem and proof:
lemma sDomain_eq_image_of_upper_span (i : Fin r) (h_i : i < ℓ + R_rate) :
    let V_i := Submodule.span 𝔽q (Set.range (sBasis β h_ℓ_add_R_rate i h_i))
    let W_i_map := polyEvalLinearMap (normalizedW 𝔽q β i)
      (normalizedW_is_additive 𝔽q β i)
    sDomain 𝔽q β h_ℓ_add_R_rate i
    = Submodule.map W_i_map V_i
:= Submodule.span 𝔽q (Set.range (sBasis β h_ℓ_add_R_rate i h_i))
    let W_i_map := polyEvalLinearMap (normalizedW 𝔽q β i)
      (normalizedW_is_additive 𝔽q β i)
    sDomain 𝔽q β h_ℓ_add_R_rate i
    = Submodule.map W_i_map V_i :=
by
  -- Proof: U_{ℓ+R} is the direct sum of Uᵢ and Vᵢ.
  -- Any x in U_{ℓ+R} can be written as u + v where u ∈ Uᵢ and v ∈ Vᵢ.
  -- Ŵᵢ(x) = Ŵᵢ(u+v) = Ŵᵢ(u) + Ŵᵢ(v) = 0 + Ŵᵢ(v) = Ŵᵢ(v).
  -- So the image of U_{ℓ+R} is the same as the image of Vᵢ.

  -- Define V_i and W_i_map for use in the proof
  set V_i := Submodule.span 𝔽q (Set.range (sBasis β h_ℓ_add_R_rate i h_i))
  set W_i_map := polyEvalLinearMap (normalizedW 𝔽q β i)
    (normalizedW_is_additive 𝔽q β i)

  -- First, show that U_{ℓ+R} = U_i ⊔ V_i (direct sum)
  have h_span_supremum_decomposition : U 𝔽q β ⟨ℓ + R_rate, h_ℓ_add_R_rate⟩
    = U 𝔽q β i ⊔ V_i := by
    unfold U
    -- U_{ℓ+R} is the span of {β₀, ..., β_{ℓ+R-1}}
    -- U_i is the span of {β₀, ..., β_{i-1}}
    -- V_i is the span of {β_i, ..., β_{ℓ+R-1}}
    have h_ico : Set.Ico 0 ⟨ℓ + R_rate, h_ℓ_add_R_rate⟩
      = Set.Ico 0 i ∪ Set.Ico i ⟨ℓ + R_rate, h_ℓ_add_R_rate⟩ := by
      ext k
      simp only [Set.mem_Ico, Fin.zero_le, true_and, Set.mem_union]
      constructor
      · intro h
        by_cases hk : k < i
        · left; omega
        · right; exact ⟨Nat.le_of_not_lt hk, by omega⟩
      · intro h
        cases h with
        | inl h => exact Fin.lt_trans h h_i
        | inr h => exact h.2

    rw [h_ico, Set.image_union, Submodule.span_union]
    congr
    -- ⊢ β '' Set.Ico i (ℓ + R_rate)
    -- = Set.range (sBasis β (h_ℓ_add_R_rate:=h_ℓ_add_R_rate) i h_i)
    -- Now how that the image of Set.Ico i (ℓ + R_rate)
    -- (from the definition of U_{ℓ+R}) is the same as V_i
    rw [sBasis_range_eq β h_ℓ_add_R_rate i h_i]

  -- Now show that the image of U_{ℓ+R} under W_i_map is the same as the image of V_i
  rw [sDomain, h_span_supremum_decomposition, Submodule.map_sup]

  -- The image of U_i under W_i_map is {0} because W_i vanishes on U_i
  have h_U_i_image : Submodule.map W_i_map (U 𝔽q β i) = ⊥ := by
    -- Show that any element in the image is 0
    apply (Submodule.eq_bot_iff _).mpr
    intro x hx
    -- x ∈ Submodule.map W_i_map (U 𝔽q β i) means x = W_i_map(y) for some y ∈ U_i
    rcases Submodule.mem_map.mp hx with ⟨y, hy, rfl⟩
    -- Show that W_i_map y = 0 for any y ∈ U_i
    have h_eval_zero : (normalizedW 𝔽q β i).eval y = 0 :=
      normalizedWᵢ_vanishing 𝔽q β i y hy
    exact h_eval_zero

  -- Combine the results: ⊥ ⊔ V = V
  rw [h_U_i_image]
  rw [bot_sup_eq]