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

abbrev RScodeSet (domain : ι ↪ F) (deg : ℕ) : Set (ι → F) := (ReedSolomon.code domain deg).carrier

open Classical in

def toFinset (domain : ι ↪ F) (deg : ℕ) : Finset (ι → F) :=
  (RScodeSet domain deg).toFinset

end

section

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

section

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

end

open LinearCode

section

open NNReal

variable [Field F]

end

section

def constantCode {α : Type*} (x : α) (ι' : Type*) [Fintype ι'] : ι' → α := fun _ ↦ x

variable [Semiring F] {x : F} [Fintype ι] {α : ι ↪ F}

end

open Finset in

-- Repository definitions used:
content: def wt [Zero F]
  (v : ι → F) : ℕ := #{i | v i ≠ 0}

content: noncomputable def dim [Semiring F] (LC : LinearCode ι F) : ℕ :=
  Module.finrank F LC

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

content: def subUpFull (U : Matrix (Fin m) (Fin n) F) (r_reindex : Fin n → Fin m) :
  Matrix (Fin n) (Fin n) F := Matrix.submatrix U r_reindex id

content: @[simps]
def refl (pSpec : ProtocolSpec n) : Equiv pSpec pSpec where
  round_eq := rfl
  dir_eq := fun _ => rfl
  typeEquiv := fun _ => _root_.Equiv.refl _

content: def subLeftFull (U : Matrix (Fin m) (Fin n) F) (c_reindex : Fin m → Fin n) :
  Matrix (Fin m) (Fin m) F := Matrix.submatrix U id c_reindex

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: def coeff.{u} {F : Type u} [Semiring F] (f : F[X][Y]) (i j : ℕ) : F := (f.coeff j).coeff i

content: noncomputable def restrictLinear [Semiring F] (S : Finset ι) :
  (ι → F) →ₗ[F] (S → F) :=
{ toFun := fun f i => f i.1,
  map_add' := by admit /- proof elided -/

content: noncomputable def minDist (C : Set (n → R)) : ℕ :=
  sInf {d | ∃ u ∈ C, ∃ v ∈ C, u ≠ v ∧ hammingDist u v = d}

content: noncomputable def minWtCodewords [Semiring F] (LC : LinearCode ι F) : ℕ :=
  sInf {w | ∃ c ∈ LC, c ≠ 0 ∧ Code.wt c = w}

content: def length [Semiring F] (_ : LinearCode ι F) : ℕ := Fintype.card ι

content: notation "Δ" IC => minDist IC

-- Repository lemmas used:
content: lemma rank_eq_if_det_ne_zero {U : Matrix (Fin n) (Fin n) F} [IsDomain F] :
  Matrix.det U ≠ 0 → U.rank = n

content: lemma rank_eq_if_subUpFull_eq (h : n ≤ m) :
   (subUpFull U (Fin.castLE h)).rank = n  → U.rank = n

content: lemma full_row_rank_via_rank_subLeftFull (h : m ≤ n) :
   (subLeftFull U (Fin.castLE h)).rank = m → U.rank = m

content: lemma cRank_rank_conversion :
  ↑(U.rank) = U.cRank

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

content: lemma rank_eq_dim_fromColGenMat [CommRing F] {G : Matrix κ ι F} :
  G.rank = dim (fromColGenMat G)

content: theorem singletonBound [CommRing F] [StrongRankCondition F]
  (LC : LinearCode ι F) :
  dim LC ≤ length LC - Code.minDist (LC : Set (ι → F)) + 1

content: lemma dist_UB [CommRing F] {LC : LinearCode ι F} :
    Code.minDist (LC : Set (ι → F)) ≤ length LC

content: lemma dist_eq_minWtCodewords [CommRing F] {LC : LinearCode ι F} :
  Code.minDist (LC : Set (ι → F)) = minWtCodewords LC

content: lemma hammingDist_eq_wt_sub [CommRing F] {u v : ι → F} : hammingDist u v = Code.wt (u - v)

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

content: abbrev RScodeSet (domain : ι ↪ F) (deg : ℕ) : Set (ι → F) := (ReedSolomon.code domain deg).carrier

content: def toFinset (domain : ι ↪ F) (deg : ℕ) : Finset (ι → F) :=
  (RScodeSet domain deg).toFinset

content: def constantCode {α : Type*} (x : α) (ι' : Type*) [Fintype ι'] : ι' → α := fun _ ↦ x

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

content: lemma subUpFull_of_vandermonde_is_vandermonde (h : n ≤ m) :
  Matrix.vandermonde (α ∘ Fin.castLE h) =
  Matrix.subUpFull (nonsquare n α) (Fin.castLE h)

content: lemma subLeftFull_of_vandermonde_is_vandermonde (h : m ≤ n) :
  Matrix.vandermonde α = Matrix.subLeftFull (nonsquare n α) (Fin.castLE h)

content: lemma rank_nonsquare_eq_deg_of_deg_le (inj : Function.Injective α) (h : n ≤ m) :
  (Vandermonde.nonsquare (ι' := n) α).rank = n

content: lemma rank_nonsquare_eq_deg_of_ι_le (inj : Function.Injective α) (h : m ≤ n) :
  (Vandermonde.nonsquare (ι' := n) α).rank = m

content: @[simp]
lemma rank_nonsquare_rows_eq_min (inj : Function.Injective α) :
  (Vandermonde.nonsquare (ι' := n) α).rank = min m n

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

content: lemma genMatIsVandermonde [Fintype ι] [Field F] [DecidableEq F] [inst : NeZero m] {α : ι ↪ F} :
  fromColGenMat (Vandermonde.nonsquare (ι' := m) α) = ReedSolomon.code α m

content: lemma dim_eq_deg_of_le [NeZero n] (inj : Function.Injective α) (h : n ≤ m) :
  dim (ReedSolomon.code ⟨α, inj⟩ n) = n

content: @[simp]
lemma dist_le_length [DecidableEq F] (inj : Function.Injective α) :
    minDist ((ReedSolomon.code ⟨α, inj⟩ n) : Set (Fin m → F)) ≤ m

content: lemma card_le_card_of_count_inj {α β : Type*} [DecidableEq α] [DecidableEq β]
    {s : Multiset α} {s' : Multiset β}
  {f : α → β} (inj : Function.Injective f) (h : ∀ a : α, s.count a ≤ s'.count (f a)) :
  s.card ≤ s'.card

-- Target theorem and proof:
theorem minDist [Field F] [DecidableEq F] (inj : Function.Injective α) [NeZero n] (h : n ≤ m) :
  minDist ((ReedSolomon.code ⟨α, inj⟩ n) : Set (Fin m → F)) = m - n + 1
:= by
  have : NeZero m := by constructor; aesop
  refine le_antisymm ?p₁ ?p₂
  case p₁ =>
    have distUB := singletonBound (LC := ReedSolomon.code ⟨α, inj⟩ n)
    rw [dim_eq_deg_of_le inj h] at distUB
    simp at distUB
    zify [dist_le_length] at distUB
    omega
  case p₂ =>
    rw [dist_eq_minWtCodewords]
    apply le_csInf (by use m, constantCode 1 _; simp)
    intro b ⟨msg, ⟨p, p_deg, p_eval_on_α_eq_msg⟩, msg_neq_0, wt_c_eq_b⟩
    let zeroes : Finset _ := {i | msg i = 0}
    have eq₁ : zeroes.val.Nodup := by
      aesop (add simp [Multiset.nodup_iff_count_eq_one, Multiset.count_filter])
    have msg_zeros_lt_deg : #zeroes < n := by
      apply lt_of_le_of_lt (b := p.roots.card)
                           (hbc := lt_of_le_of_lt (Polynomial.card_roots' _)
                                                  (natDegree_lt_of_mem_degreeLT p_deg))
      exact card_le_card_of_count_inj inj fun i ↦
        if h : msg i = 0
        then suffices 0 < Multiset.count (α i) p.roots by
                rwa [@Multiset.count_eq_one_of_mem (d := eq₁) (h := by simpa [zeroes])]
              by aesop
        else by simp [zeroes, h]
    have : #zeroes + wt msg = m := by
      rw [wt, filter_card_add_filter_neg_card_eq_card]
      simp
    omega