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

-- 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: omit [DecidableEq L] [Fintype 𝔽q] hβ_lin_indep in
lemma W_ne_zero (i : Fin r) : (W 𝔽q β i) ≠ 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] 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

-- Target theorem and proof:
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
:= by
  rw [←Polynomial.count_roots]
  set f := fun c: 𝔽q => (W 𝔽q β i).comp (X - C (c • β i)) with hf
  -- ⊢ Multiset.count a (univ.prod f).roots = if a ∈ ↑(U 𝔽q β (i + 1)) then 1 else 0
  have h_prod_ne_zero: univ.prod f ≠ 0 := Prod_W_comp_X_sub_C_ne_zero 𝔽q β i
  rw [roots_prod (f := f) (s := univ (α := 𝔽q)) h_prod_ne_zero]
  set roots_f := fun c: 𝔽q => (f c).roots with hroots_f
  rw [Multiset.count_bind]
  -- ⊢ (Multiset.map (fun b ↦ Multiset.count a (roots_f b)) univ.val).sum
  -- = if a ∈ ↑(U 𝔽q β (i + 1)) then 1 else 0
  have h_roots_f_eq_roots_W : ∀ b : 𝔽q,
    roots_f b = (W 𝔽q β i).roots.map (fun r => r + (b • β i)) := by
    intro b
    rw [hroots_f, hf]
    exact roots_comp_X_sub_C (p := (W 𝔽q β i)) (a := (b • β i))
  simp_rw [h_roots_f_eq_roots_W]

  set shift_up := fun x: 𝔽q => fun r: L => r + x • β i with hshift_up
  have h_shift_up_all: ∀ x: 𝔽q, ∀ r: L, shift_up x r = r + x • β i := by
    intro x r
    rw [hshift_up]
  simp only [sum_map_val, SetLike.mem_coe]
  have h_a: ∀ x: 𝔽q, a = shift_up x (a - x • β i) := by
    intro x
    rw [hshift_up]
    simp_all only [ne_eq, implies_true, sub_add_cancel, f, roots_f, shift_up]
  conv_lhs =>
    enter [2, x] -- focus on the inner Multiset.count
    rw [h_a x]
    enter [2]
    enter [1]
    enter [r]
    rw [←h_shift_up_all x r] -- rewrite to another notation
  -- ⊢ ∑ x, Multiset.count (shift_up x (a - x • β i)) (Multiset.map (shift_up x) (W 𝔽q β i).roots)
  -- = if a ∈ ↑(U 𝔽q β (i + 1)) then 1 else 0
  have h_shift_up_inj: ∀ x: 𝔽q, Function.Injective (shift_up x) := by
    intro x
    unfold shift_up
    exact add_left_injective (x • β i)
  have h_count_map: ∀ x: 𝔽q,
    Multiset.count (shift_up x (a - x • β i)) (Multiset.map (shift_up x) (W 𝔽q β i).roots) =
    Multiset.count (a - x • β i) (W 𝔽q β i).roots := by
    -- transform to counting (a - x • β i) in the roots of Wᵢ
    intro x
    have h_shift_up_inj_x: Function.Injective (shift_up x) := h_shift_up_inj x
    simp only [Multiset.count_map_eq_count' (hf := h_shift_up_inj_x), count_roots]
  conv_lhs =>
    enter [2, x]
    rw [h_count_map x]
  -- ⊢ ∑ x, Multiset.count (a - x • β i) (W 𝔽q β i).roots
  -- = if a ∈ ↑(U 𝔽q β (i + 1)) then 1 else 0
  have h_root_lift_down := root_U_lift_down 𝔽q β i h_i_add_1 a
  have h_root_lift_up := root_U_lift_up 𝔽q β i h_i_add_1 a
  conv_lhs =>
    enter [2, x]
    simp only [count_roots]
    rw [rootMultiplicity_W]
  by_cases h_a_mem_U_i : a ∈ ↑(U 𝔽q β (i + 1))
  · -- ⊢ (∑ x, if a - x • β i ∈ ↑(U 𝔽q β i) then 1 else 0)
    -- = if a ∈ ↑(U 𝔽q β (i + 1)) then 1 else 0
    have h_true: (a ∈ ↑(U 𝔽q β (i + 1))) = True := by simp only [h_a_mem_U_i]
    rcases h_root_lift_down h_a_mem_U_i with ⟨x0, hx0, hx0_unique⟩
    conv =>
      rhs
      -- | if a ∈ ↑(U 𝔽q β (i + 1)) then 1 else 0 => reduce this to 1
      enter [1]
      exact h_true -- maybe there can be a better way to do this
    rw [ite_true]
    classical
    -- ⊢ (∑ x, if a - x • β i ∈ ↑(U 𝔽q β i) then 1 else 0) = 1
    have h_true: ∀ x: 𝔽q,
      if x = x0 then a - x • β i ∈ ↑(U 𝔽q β i) else a - x • β i ∉ ↑(U 𝔽q β i) := by
      intro x
      by_cases h_x_eq_x0 : x = x0
      · rw [if_pos h_x_eq_x0] -- ⊢ a - x • β i ∈ U 𝔽q β i
        rw [←h_x_eq_x0] at hx0
        exact hx0
      · rw [if_neg h_x_eq_x0] -- ⊢ a - x • β i ∉ U 𝔽q β i
        by_contra h_mem
        have h1 := hx0_unique x
        simp only [h_mem, forall_const] at h1
        contradiction

    have h_true_x: ∀ x: 𝔽q, (a - x • β i ∈ ↑(U 𝔽q β i)) = if x = x0 then True else False := by
      intro x
      by_cases h_x_eq_x0 : x = x0
      · rw [if_pos h_x_eq_x0]
        rw [←h_x_eq_x0] at hx0
        simp only [hx0]
      · rw [if_neg h_x_eq_x0]
        by_contra h_mem
        push_neg at h_mem
        simp only [ne_eq, eq_iff_iff, iff_false, not_not] at h_mem
        have h2 := hx0_unique x
        simp only [h_mem, forall_const] at h2
        contradiction
    conv =>
      lhs
      enter [2, x]
      simp only [SetLike.mem_coe, h_true_x x, if_false_right, and_true]
    rw [sum_ite_eq']
    simp only [mem_univ, ↓reduceIte]
  · -- ⊢ (∑ x, if a - x • β i ∈ ↑(U 𝔽q β i) then 1 else 0)
    -- = if a ∈ ↑(U 𝔽q β (i + 1)) then 1 else 0
    have h_false: (a ∈ ↑(U 𝔽q β (i + 1))) = False := by simp only [h_a_mem_U_i]
    conv =>
      rhs -- | if a ∈ ↑(U 𝔽q β (i + 1)) then 1 else 0 => reduce this to 1
      enter [1]
      exact h_false -- maybe there can be a better way to do this
    rw [ite_false]

    have h_zero_x: ∀ x: 𝔽q, (a - x • β i ∈ ↑(U 𝔽q β i)) = False := by
      intro x
      by_contra h_mem
      simp only [eq_iff_iff, iff_false, not_not] at h_mem -- h_mem : a - x • β i ∈ U 𝔽q β i
      have h_a_mem_U_i := h_root_lift_up x h_mem
      contradiction

    conv =>
      lhs
      enter [2, x]
      simp only [SetLike.mem_coe, h_zero_x x, if_false_right, and_true]
    simp only [↓reduceIte, sum_const_zero]