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

content: theorem prod_poly_sub_C_eq_poly_pow_card_sub_poly_in_L
  (p : L[X]) :
  (∏ c ∈ (Finset.univ : Finset Fq), (p - Polynomial.C (algebraMap Fq L c))) =
    p^(Fintype.card Fq) - p

content: theorem prod_X_sub_C_eq_X_pow_card_sub_X_in_L :
  (∏ c ∈ (Finset.univ : Finset Fq), (Polynomial.X - Polynomial.C (algebraMap Fq L c))) =
    Polynomial.X^(Fintype.card Fq) - Polynomial.X

content: theorem prod_X_sub_C_eq_X_pow_card_sub_X :
  (∏ c ∈ (Finset.univ : Finset Fq), (Polynomial.X - Polynomial.C c)) =
    Polynomial.X^(Fintype.card Fq) - Polynomial.X

content: theorem linear_map_of_comp_to_linear_map_of_eval (f : L[X])
  (h_f_linear : IsLinearMap (R := Fq) (M := L[X]) (M₂ := L[X])
    (f := fun inner_p ↦ f.comp inner_p)) :
    IsLinearMap (R := Fq) (M := L) (M₂ := L) (f := fun x ↦ f.eval x)

-- 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ᵢ_eval_βᵢ_neq_zero
    (i : Fin r): (W 𝔽q β i).eval (β 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

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

-- Target theorem and proof:
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)
:= by
  intro p
  set W_i := W 𝔽q β i
  set q := Fintype.card 𝔽q
  set v := W_i.eval (β i)

  -- First, we must prove that v is non-zero to use its inverse.
  have hv_ne_zero : v ≠ 0 := by
    unfold v W_i
    exact Wᵢ_eval_βᵢ_neq_zero 𝔽q β i

  -- Proof flow:
  -- `Wᵢ₊₁(X) = ∏_{c ∈ 𝔽q} (Wᵢ ∘ (X - c • βᵢ))` -- from W_prod_comp_decomposition
    -- `= ∏_{c ∈ 𝔽q} (Wᵢ(X) - c • Wᵢ(βᵢ))` -- linearity of Wᵢ
    -- `= ∏_{c ∈ 𝔽q} (Wᵢ(X) - c • v)`
    -- `= v² ∏_{c ∈ 𝔽q} (v⁻¹ • Wᵢ(X) - c)`
    -- `= v² (v⁻² • Wᵢ(X)² - v⁻¹ • Wᵢ(X))` => FLT (prod_X_sub_C_eq_X_pow_card_sub_X_in_L)
    -- `= Wᵢ(X)² - v • Wᵢ(X)` => Q.E.D

  have h_scalar_smul_eq_C_v_mul: ∀ s: L, ∀ p: L[X], s • p = C s * p := by
    intro s p
    exact smul_eq_C_mul s
  have h_v_smul_v_inv_eq_one: v • v⁻¹ = 1 := by
    simp only [smul_eq_mul]
    exact CommGroupWithZero.mul_inv_cancel v hv_ne_zero
  have h_v_mul_v_inv_eq_one: v * v⁻¹ = 1 := by
    exact h_v_smul_v_inv_eq_one
  -- The main proof using a chain of equalities (the `calc` block).
  calc
    (W 𝔽q β (i + 1)).comp p
    _ = (∏ c: 𝔽q, (W_i).comp (X - C (c • β i))).comp p := by
      have h_res := W_prod_comp_decomposition 𝔽q β (i+1) (by
        apply Fin.mk_lt_of_lt_val
        rw [Fin.val_add_one' (a := i) (h_a_add_1 := h_i_add_1), Nat.zero_mod]
        omega
      )
      rw [h_res]
      simp only [add_sub_cancel_right]
      rfl
    -- Step 2: Apply the linearity property of Wᵢ as a polynomial.
    _ = (∏ c: 𝔽q, (W_i - C (W_i.eval (c • β i)))).comp p := by
      congr
      funext c
      -- We apply the transformation inside the product for each element `c`.
      -- apply Finset.prod_congr rfl
      -- ⊢ W_i.comp (X - C (c • β i)) = W_i - C (eval (c • β i) W_i)
      exact comp_sub_C_of_linear_eval (p := W_i) (h_lin := h_prev_linear_map) (a := (c • β i))
    -- Step 3: Apply the linearity of Wᵢ's *evaluation map* to the constant term.
    -- Hypothesis: `h_prev_linear_map.map_smul`
    _ = (∏ c: 𝔽q, (W_i - C (c • v))).comp p := by
      congr
      funext c
      -- ⊢ W_i - C (eval (c • β i) W_i) = W_i - C (c • v)
      congr
      -- ⊢ eval (c • β i) W_i = c • v
      -- Use the linearity of the evaluation map, not the composition map
      have h_eval_linear := Polynomial.linear_map_of_comp_to_linear_map_of_eval (f := (W 𝔽q β i))
        (h_f_linear := h_prev_linear_map)
      exact h_eval_linear.map_smul c (β i)
    -- Step 4: Perform the final algebraic transformation.
    _ = (C (v^q) * (∏ c: 𝔽q, (C (v⁻¹) * W_i - C (algebraMap 𝔽q L c)))).comp p := by
      congr
      calc
        _ = ∏ c: 𝔽q, (v • (v⁻¹ • W_i - C (algebraMap 𝔽q L c))) := by
          apply Finset.prod_congr rfl
          intro c _
          rw [smul_sub]
          -- ⊢ W_i - C (c • v) = v • v⁻¹ • W_i - v • C ((algebraMap 𝔽q L) c)
          rw [smul_C, smul_eq_mul, map_mul]
          rw [←smul_assoc]
          rw [h_v_smul_v_inv_eq_one]
          rw [one_smul]
          rw [sub_right_inj]
          -- ⊢ C (c • v) = C v * C ((algebraMap 𝔽q L) c)
          rw [←C_mul]
          -- ⊢ C (c • v) = C (v * (algebraMap 𝔽q L) c)
          have h_c_smul_v: c • v = (algebraMap 𝔽q L c) • v := by
            exact algebra_compatible_smul L c v
          rw [h_c_smul_v]
          rw [mul_comm]
          rw [smul_eq_mul]
        _ = ∏ c: 𝔽q, (C v * (v⁻¹ • W_i - C (algebraMap 𝔽q L c))) := by
          apply Finset.prod_congr rfl
          intro c _
          rw [h_scalar_smul_eq_C_v_mul]
        _ = C (v^q) * (∏ c: 𝔽q, (C v⁻¹ * W_i - C (algebraMap 𝔽q L c))) := by
          -- rw [Finset.prod_mul_distrib]
          -- rw [Finset.prod_const, Finset.card_univ]
          rw [Finset.prod_mul_distrib]
          conv_lhs =>
            enter [2]
            enter [2]
            rw [h_scalar_smul_eq_C_v_mul]
          congr
          -- ⊢ ∏ (x: 𝔽q), C v = C (v ^ q)
          rw [Finset.prod_const, Finset.card_univ]
          unfold q
          exact Eq.symm C_pow
    _ = (C (v^q) * ((C v⁻¹ * W_i)^q - (C v⁻¹ * W_i))).comp p := by
      congr
      -- ⊢ ∏ c, (C v⁻¹ * W_i - C ((algebraMap 𝔽q L) c)) = (C v⁻¹ * W_i) ^ q - C v⁻¹ * W_i
      rw [Polynomial.prod_poly_sub_C_eq_poly_pow_card_sub_poly_in_L (p := C v⁻¹ * W_i)]
    _ = (C (v^q) * C (v⁻¹^q) * W_i^q - C (v^q) * C v⁻¹ * W_i).comp p := by
      congr
      rw [mul_sub]
      conv_lhs =>
        rw [mul_pow, ←mul_assoc, ←mul_assoc, ←C_pow]
    _ = (W_i^q - C (v^(q-1)) * W_i).comp p := by
      congr
      · rw [←C_mul, ←mul_pow, h_v_mul_v_inv_eq_one, one_pow, C_1, one_mul]
      · rw [←C_mul]
        have h_v_pow_q_minus_1: v^q * v⁻¹ = v^(q-1) := by
          rw [pow_sub₀ (a := v) (m := q) (n := 1) (ha := hv_ne_zero) (h := by exact NeZero.one_le)]
          -- ⊢ v ^ q * v⁻¹ = v ^ q * (v ^ 1)⁻¹
          congr
          norm_num
        rw [h_v_pow_q_minus_1]
    _ = (W_i^q - C (eval (β i) W_i) ^ (q - 1) * W_i).comp p := by
      simp only [map_pow, W_i, q, v]
    _ = (W_i^q).comp p - (C (eval (β i) W_i) ^ (q - 1) * W_i).comp p := by
      rw [sub_comp]
    _ = (W_i.comp p)^q - (C (eval (β i) W_i) ^ (q - 1)) * (W_i.comp p) := by
      rw [pow_comp, mul_comp]
      conv_lhs =>
        rw [pow_comp]
        rw [C_comp (a := (eval (β i) W_i)) (p := p)]