-- ArkLib/Data/CodingTheory/ReedSolomon.lean

import ArkLib.Data.MvPolynomial.LinearMvExtension

import ArkLib.Data.Polynomial.Interface

import Mathlib.LinearAlgebra.Lagrange

import Mathlib.RingTheory.Henselian

namespace ReedSolomon

open Polynomial NNReal

variable {F : Type*} {ι : Type*} (domain : ι ↪ F)

def evalOnPoints [Semiring F] : F[X] →ₗ[F] (ι → F) where
  toFun := fun p => fun x => p.eval (domain x)
  map_add' := fun x y => by admit /- proof elided -/

def code (deg : ℕ) [Semiring F]: Submodule F (ι → F) :=
  (Polynomial.degreeLT F deg).map (evalOnPoints domain)

variable [Semiring F]

end ReedSolomon

open Polynomial Matrix Code LinearCode

variable {F ι ι' : Type*}
         {C : Set (ι → F)}

noncomputable section

namespace Vandermonde

def nonsquare [Semiring F] (ι' : ℕ) (α : ι → F) : Matrix ι (Fin ι') F :=
  Matrix.of fun i j => (α i) ^ j.1

section

variable [CommRing F] {m n : ℕ} {α : Fin m → F}

section

variable [IsDomain F]

end

end

end Vandermonde

namespace ReedSolomonCode

section

open Finset Function

open scoped BigOperators

variable {ι : Type*} [Fintype ι] [Nonempty ι]
         {F : Type*} [Field F] [Fintype F]

open Classical in

end

section

variable {deg m n : ℕ} {α : Fin m → F}

section

variable [Semiring F] {p : F[X]}

end

open LinearCode

-- Repository definitions used:
content: def polynomialOfCoeffs (coeffs : Fin deg → F) : F[X] :=
  ⟨
    Finset.map ⟨Fin.val, Fin.val_injective⟩ {i | coeffs i ≠ 0},
    fun i ↦ if h : i < deg then coeffs ⟨i, h⟩ else 0,
    fun a ↦ by admit /- proof elided -/
  ⟩

content: def liftF' (f : ℕ → α) : Fin n → α :=
  fun m ↦ f m.1

content: noncomputable def fromColGenMat [CommRing F] (G : Matrix ι κ F) : LinearCode ι F :=
  LinearMap.range G.mulVecLin

content: abbrev LinearCode.{u, v} (ι : Type u) [Fintype ι] (F : Type v) [Semiring F] : Type (max u v) :=
  Submodule F (ι → F)

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

-- Repository lemmas used:
content: @[simp]
lemma liftF'_p_coeff {p : F[X]} {k : ℕ} {i : Fin k} : liftF' p.coeff i = p.coeff i

content: @[simp]
lemma coeff_polynomialOfCoeffs_eq_coeffs :
  Fin.liftF' (polynomialOfCoeffs coeffs).coeff = coeffs

-- Local definitions used:
content: def evalOnPoints [Semiring F] : F[X] →ₗ[F] (ι → F) where
  toFun := fun p => fun x => p.eval (domain x)
  map_add' := fun x y => by admit /- proof elided -/

content: def code (deg : ℕ) [Semiring F]: Submodule F (ι → F) :=
  (Polynomial.degreeLT F deg).map (evalOnPoints domain)

content: def nonsquare [Semiring F] (ι' : ℕ) (α : ι → F) : Matrix ι (Fin ι') F :=
  Matrix.of fun i j => (α i) ^ j.1

-- Local lemmas used:
content: lemma nonsquare_mulVecLin [CommSemiring F] {ι' : ℕ} {α₁ : ι ↪ F} {α₂ : Fin ι' → F} {i : ι} :
  (nonsquare ι' α₁).mulVecLin α₂ i = ∑ x, α₂ x * α₁ i ^ x.1

content: theorem mulVecLin_coeff_vandermondens_eq_eval_matrixOfPolynomials
  {n : ℕ} [NeZero n] {v : ι ↪ F} {p : F[X]} (h_deg : p.natDegree < n) :
  (Vandermonde.nonsquare (ι' := n) v).mulVecLin (Fin.liftF' p.coeff) =
  fun i => p.eval (v i)

content: lemma natDegree_lt_of_mem_degreeLT [NeZero deg] (h : p ∈ degreeLT F deg) : p.natDegree < deg

-- Target theorem and proof:
lemma genMatIsVandermonde [Fintype ι] [Field F] [DecidableEq F] [inst : NeZero m] {α : ι ↪ F} :
  fromColGenMat (Vandermonde.nonsquare (ι' := m) α) = ReedSolomon.code α m
:= by
  unfold fromColGenMat ReedSolomon.code
  ext x; rw [LinearMap.mem_range, Submodule.mem_map]
  refine ⟨
    fun ⟨coeffs, h⟩ ↦ ⟨polynomialOfCoeffs coeffs, h.symm ▸ ?p₁⟩,
    fun ⟨p, h⟩ ↦ ⟨Fin.liftF' p.coeff, ?p₂⟩
  ⟩
  · rw [
      ←coeff_polynomialOfCoeffs_eq_coeffs (coeffs := coeffs),
      Vandermonde.mulVecLin_coeff_vandermondens_eq_eval_matrixOfPolynomials (by simp)
    ]
    simp [ReedSolomon.evalOnPoints]
  · exact h.2 ▸ Vandermonde.mulVecLin_coeff_vandermondens_eq_eval_matrixOfPolynomials
                  (natDegree_lt_of_mem_degreeLT h.1)