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

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

end LinearityOfSubspaceVanishingPolynomials

section NovelPolynomialBasisProof

noncomputable def basisVectors (ℓ : Nat) (h_ℓ : ℓ ≤ r) :
  Fin (2 ^ ℓ) → L⦃<2^ℓ⦄[X] :=
  fun j => ⟨Xⱼ 𝔽q β ℓ h_ℓ j, by admit /- proof elided -/
  ⟩

abbrev CoeffVecSpace (L : Type u) (ℓ : Nat) := Fin (2^ℓ) → L

def toCoeffsVec (ℓ : Nat) : L⦃<2^ℓ⦄[X] →ₗ[L] CoeffVecSpace L ℓ where
  toFun := fun p => fun i => p.val.coeff i.val
  map_add' := fun p q => by admit /- proof elided -/

noncomputable def changeOfBasisMatrix (ℓ : Nat) (h_ℓ : ℓ ≤ r) : Matrix (Fin (2^ℓ)) (Fin (2^ℓ)) L :=
    fun j i => (toCoeffsVec (L := L) (ℓ := ℓ) (
      basisVectors 𝔽q β ℓ h_ℓ j)) i

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

content: instance finiteDimensional_degreeLT {n : ℕ} (h_n_pos : 0 < n) :
  FiniteDimensional L L⦃< n⦄[X] :=

content: def coeff.{u} {F : Type u} [Semiring F] (f : F[X][Y]) (i j : ℕ) : F := (f.coeff j).coeff i

content: noncomputable def monomialBasisOfDegreeLT {n : ℕ} : Basis (Fin n) L (L⦃< n⦄[X]) :=

-- Repository lemmas used:
content: theorem getBit_repr {ℓ : Nat} : ∀ j, j < 2^ℓ →
  j = ∑ k ∈ Finset.Icc 0 (ℓ-1), (getBit k j) * 2^k

content: lemma getBit_of_shiftRight {n p : ℕ}:
  ∀ k, getBit k (n >>> p) = getBit (k+p) n

content: lemma getBit_eq_zero_or_one {k n : Nat} :
  getBit k n = 0 ∨ getBit k n = 1

content: theorem finrank_degreeLT_n (n : ℕ) : Module.finrank L (L⦃< n⦄[X]) = n

-- 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: noncomputable def normalizedW (i : Fin r) : L[X] :=
  C (1 / (W 𝔽q β i).eval (β i)) * W 𝔽q β i

content: noncomputable def Xⱼ (ℓ : ℕ) (h_ℓ : ℓ ≤ r) (j : Fin (2 ^ ℓ)) : L[X] :=
  (Finset.univ : Finset (Fin ℓ)).prod
    (fun i => (normalizedW 𝔽q β (Fin.castLE h_ℓ i))^(Nat.getBit i j))

content: noncomputable def basisVectors (ℓ : Nat) (h_ℓ : ℓ ≤ r) :
  Fin (2 ^ ℓ) → L⦃<2^ℓ⦄[X] :=
  fun j => ⟨Xⱼ 𝔽q β ℓ h_ℓ j, by admit /- proof elided -/
  ⟩

content: abbrev CoeffVecSpace (L : Type u) (ℓ : Nat) := Fin (2^ℓ) → L

content: def toCoeffsVec (ℓ : Nat) : L⦃<2^ℓ⦄[X] →ₗ[L] CoeffVecSpace L ℓ where
  toFun := fun p => fun i => p.val.coeff i.val
  map_add' := fun p q => by admit /- proof elided -/

content: noncomputable def changeOfBasisMatrix (ℓ : Nat) (h_ℓ : ℓ ≤ r) : Matrix (Fin (2^ℓ)) (Fin (2^ℓ)) L :=
    fun j i => (toCoeffsVec (L := L) (ℓ := ℓ) (
      basisVectors 𝔽q β ℓ h_ℓ j)) i

-- Local lemmas used:
content: omit [Fintype L] [Fintype 𝔽q] h_Fq_char_prime in
lemma finrank_U (i : Fin r) :
  Module.finrank 𝔽q (U 𝔽q β i) = i

content: lemma U_card (i : Fin r) :
    Fintype.card (U 𝔽q β i) = (Fintype.card 𝔽q)^i.val

content: lemma βᵢ_not_in_Uᵢ (i : Fin r) :
    β i ∉ U 𝔽q β i

content: lemma degree_W (i : Fin r) : (W 𝔽q β i).degree = (Fintype.card 𝔽q)^i.val

content: lemma Wᵢ_eval_βᵢ_neq_zero
    (i : Fin r): (W 𝔽q β i).eval (β i) ≠ 0

content: lemma degree_normalizedW (i : Fin r) :
  (normalizedW 𝔽q β i).degree = (Fintype.card 𝔽q)^(i.val)

content: lemma Xⱼ_zero_eq_one (ℓ : ℕ) (h_ℓ : ℓ ≤ r) :
  Xⱼ 𝔽q β ℓ h_ℓ ⟨0, by exact Nat.two_pow_pos ℓ⟩ = 1

content: lemma degree_Xⱼ (ℓ : ℕ) (h_ℓ : ℓ ≤ r) (j : Fin (2 ^ ℓ)) :
  (Xⱼ 𝔽q β ℓ h_ℓ j).degree = j

content: lemma linearIndependent_rows_of_lower_triangular_ne_zero_diag
  {n : ℕ} {R : Type*} [Field R] (A : Matrix (Fin n) (Fin n) R)
  (h_lower_triangular : A.BlockTriangular ⇑OrderDual.toDual) (h_diag : ∀ i, A i i ≠ 0) :
  LinearIndependent R A

content: omit h_Fq_char_prime in
theorem changeOfBasisMatrix_lower_triangular
  (ℓ : Nat) (h_ℓ : ℓ ≤ r) :
  (changeOfBasisMatrix 𝔽q β ℓ h_ℓ).BlockTriangular ⇑OrderDual.toDual

content: omit h_Fq_char_prime in
theorem changeOfBasisMatrix_diag_ne_zero
  (ℓ : Nat) (h_ℓ : ℓ ≤ r) :
  (∀ i, (changeOfBasisMatrix 𝔽q β ℓ h_ℓ) i i ≠ 0)

content: lemma coeff_vectors_linear_independent
  (ℓ : Nat) (h_ℓ : ℓ ≤ r) :
    LinearIndependent L (toCoeffsVec (ℓ := ℓ) ∘ (basisVectors 𝔽q β ℓ h_ℓ))

content: theorem basisVectors_linear_independent (ℓ : Nat) (h_ℓ : ℓ ≤ r) :
    LinearIndependent L (basisVectors 𝔽q β ℓ h_ℓ)

-- Target theorem and proof:
theorem basisVectors_span (ℓ : Nat) (h_ℓ : ℓ ≤ r) :
    Submodule.span L (Set.range (basisVectors 𝔽q β ℓ h_ℓ)) = ⊤
:= by
  have h_li := basisVectors_linear_independent 𝔽q β ℓ h_ℓ
  let n := 2 ^ ℓ
  have h_n: n = 2 ^ ℓ := by omega
  have h_n_pos: 0 < n := by
    rw [h_n]
    exact Nat.two_pow_pos ℓ
  have h_finrank_eq_n : Module.finrank L (L⦃< n⦄[X]) = n := finrank_degreeLT_n n
  -- We have `n` linearly independent vectors in an `n`-dimensional space.
  -- The dimension of their span is `n`.
  have h_span_finrank : Module.finrank L (Submodule.span L (Set.range (
    basisVectors 𝔽q β ℓ h_ℓ))) = n := by
    rw [finrank_span_eq_card h_li, Fintype.card_fin]
  -- A subspace with the same dimension as the ambient space must be the whole space.
  rw [←h_finrank_eq_n] at h_span_finrank
  have inst_finite_dim : FiniteDimensional (K := L) (V := L⦃< n⦄[X]) :=
    finiteDimensional_degreeLT (h_n_pos := by omega)
  apply Submodule.eq_top_of_finrank_eq (K := L) (V := L⦃< n⦄[X])
  exact h_span_finrank