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

import ArkLib.Data.Nat.Bitwise

import ArkLib.Data.Polynomial.Frobenius

import ArkLib.Data.Polynomial.MonomialBasis

import Mathlib.LinearAlgebra.StdBasis

import Mathlib.Algebra.Polynomial.Degree.Definitions

open Polynomial FiniteDimensional Finset 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]
  [h_Fq_char_prime : Fact (Nat.Prime (ringChar 𝔽q))] [hF₂ : Fact (Fintype.card 𝔽q = 2)]

variable [Algebra 𝔽q L]

variable (h_dim : Module.finrank 𝔽q L = r)

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

section LinearSubspaces

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

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

end LinearSubspaces

section LinearityOfSubspaceVanishingPolynomials

@[simps!]
noncomputable def algEquivAevalXSubC {R : Type*} [CommRing R] (t : R) : R[X] ≃ₐ[R] R[X] :=

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

-- Repository lemmas used:
content: lemma Fin.lt_succ' (a : Fin r) (h_a_add_1 : a + 1 < r) : a < a + 1

content: lemma Fin.val_add_one' (a : Fin r) (h_a_add_1 : a + 1 < r) : (a + 1).val = a.val + 1

content: lemma Fin.le_succ (a : Fin r) (h_a_add_1 : a + 1 < r) : a ≤ a + 1

content: lemma Fin.le_iff_lt_succ (a b : Fin r) (h_b : b + 1 < r) : a ≤ b ↔ a < b + 1

content: lemma Fin.val_sub_one (a : Fin r) (h_a_sub_1 : a > 0) : (a - 1).val = a.val - 1

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

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

content: @[simps!]
noncomputable def algEquivAevalXSubC {R : Type*} [CommRing R] (t : R) : R[X] ≃ₐ[R] R[X] :=

-- Local lemmas used:
content: lemma βᵢ_not_in_Uᵢ (i : Fin r) :
    β i ∉ U 𝔽q β i

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: 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: lemma W_monic (i : Fin r) : (W 𝔽q β i).Monic

content: omit [DecidableEq L] [Fintype 𝔽q] hβ_lin_indep in
lemma W_ne_zero (i : Fin r) : (W 𝔽q β i) ≠ 0

content: lemma W_splits (i : Fin r) : (W 𝔽q β i).Splits (RingHom.id L)

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] 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_comp_X_sub_C (p : L[X]) (a x : L) :
    rootMultiplicity x (p.comp (X - C a)) = rootMultiplicity (x - a) p

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

-- Target theorem and proof:
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)))
:= by
  -- ⊢ W 𝔽q β i = ∏ c, (W 𝔽q β (i - 1)).comp (X - C (c • β (i - 1)))
  -- Define P and Q for clarity
  set P := W 𝔽q β i
  set Q := ∏ c: 𝔽q, (W 𝔽q β (i-1)).comp (X - C (c • β (i-1)))

-- c : 𝔽q => univ
-- c ∈ finsetX

  -- STRATEGY: Prove P = Q by showing they are monic, split, and have the same roots.

  -- 1. Show P and Q are MONIC.
  have hP_monic : P.Monic := W_monic (𝔽q := 𝔽q) (β := β) (i :=i)
  have hQ_monic : Q.Monic := by
    apply Polynomial.monic_prod_of_monic; intro c _
    apply Monic.comp
    · exact W_monic (𝔽q := 𝔽q) (β := β) (i :=(i-1))
    · -- ⊢ (X - C (c • β (i - 1))).Monic
      exact Polynomial.monic_X_sub_C (c • β (i - 1))
    · conv_lhs => rw [natDegree_sub_C, natDegree_X]
      norm_num
  -- 2. Show P and Q SPLIT over L.
  have hP_splits : P.Splits (RingHom.id L) := W_splits 𝔽q β i
  have hQ_splits : Q.Splits (RingHom.id L) := by
    apply Polynomial.splits_prod
    intro c _
    -- Composition of a splitting polynomial with a linear polynomial also splits.
    -- ⊢ Splits (RingHom.id L) ((W 𝔽q β (i - 1)).comp (X - C (c • β (i - 1))))
    apply Splits.comp_of_degree_le_one
    · exact degree_X_sub_C_le (c • β (i - 1))
    · -- ⊢ Splits (RingHom.id L) (W 𝔽q β (i - 1))
      exact W_splits 𝔽q β (i-1)

  -- 3. Show P and Q have the same ROOTS.
  have h_roots_eq : P.roots = Q.roots := by
    -- First, characterize the roots of P. They are the elements of Uᵢ.
    unfold P Q
    ext u
    rw [Polynomial.count_roots, Polynomial.count_roots]
    rw [rootMultiplicity_W]
    conv_rhs =>
      rw [rootMultiplicity_prod_W_comp_X_sub_C 𝔽q β (h_i_add_1 := by
        rw [Fin.val_sub_one (a := i) (h_a_sub_1 := by omega)]
        omega
      )]
    -- ⊢ (if u ∈ ↑(U 𝔽q β i) then 1 else 0) = if u ∈ ↑(U 𝔽q β (i - 1 + 1)) then 1 else 0
    have h_i : i - 1 + 1 = i := by simp only [sub_add_cancel]
    rw [h_i]

  -- 4. CONCLUSION: Since P and Q are monic, split, and have the same roots, they are equal.
  have hP_eq_prod := Polynomial.eq_prod_roots_of_monic_of_splits_id hP_monic hP_splits
  have hQ_eq_prod := Polynomial.eq_prod_roots_of_monic_of_splits_id hQ_monic hQ_splits
  rw [hP_eq_prod, hQ_eq_prod, h_roots_eq]