-- ArkLib/Data/CodingTheory/BerlekampWelch/ElocPoly.lean

import Init.Data.List.FinRange

import Mathlib.Algebra.Field.Basic

import Mathlib.Algebra.Polynomial.Basic

import Mathlib.Algebra.Polynomial.Degree.Definitions

import Mathlib.Algebra.Polynomial.FieldDivision

import Mathlib.Data.Finset.Insert

import Mathlib.Data.Fintype.Card

import Mathlib.Data.Matrix.Mul

import ArkLib.Data.CodingTheory.Basic

import ArkLib.Data.Fin.Lift

namespace BerlekampWelch

variable {F : Type} [Field F]
         {m n : ℕ} {p : Polynomial F}

variable [DecidableEq F]

section ElocPoly

open Polynomial

protected noncomputable def ElocPoly (n : ℕ) (ωs f : ℕ → F) (p : Polynomial F) : Polynomial F :=
  List.prod <| (List.range n).map fun i =>
    if f i = p.eval (ωs i)
    then 1
    else X - C (ωs i)

section

open BerlekampWelch (ElocPoly)

variable {ωs f : ℕ → F}

open BerlekampWelch (elocPoly_succ) in

section

open Fin

open BerlekampWelch (elocPoly_congr)

noncomputable def ElocPolyF (ωs f : Fin n → F) (p : Polynomial F) : Polynomial F :=
  ElocPoly n (liftF ωs) (liftF f) p

open BerlekampWelch
  (elocPolyF_eq_elocPoly' elocPoly_leftF_leftF_eq_contract
   elocPoly_zero elocPoly_succ)

open Fin

-- Repository definitions used:
content: notation "Δ₀(" u ", " v ")" => hammingDist u v

content: notation "Δ₀(" u ", " C ")" => distFromCode u C

content: scoped macro_rules
  | `(ρ $t:term) => `(LinearCode.rate $t)

content: def liftF (f : Fin n → α) : ℕ → α :=
  fun m ↦ if h : m < n then f ⟨m, h⟩ else 0

content: noncomputable def rate [Semiring F] (LC : LinearCode ι F) : ℚ≥0 :=
  (dim LC : ℚ≥0) / length LC

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 length [Semiring F] (_ : LinearCode ι F) : ℕ := Fintype.card ι

content: abbrev contract (m : ℕ) (f : Fin n → α) := liftF (liftF' (n := m) (liftF f))

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

-- Repository lemmas used:
content: lemma contract_eq_liftF_of_lt {k : ℕ} (h₁ : k < m) :
  contract m f' k = liftF f' k

content: @[simp]
lemma liftF_succ {f : Fin (n + 1) → α} : liftF f n = f ⟨n, Nat.lt_add_one _⟩

-- Local definitions used:
content: protected noncomputable def ElocPoly (n : ℕ) (ωs f : ℕ → F) (p : Polynomial F) : Polynomial F :=
  List.prod <| (List.range n).map fun i =>
    if f i = p.eval (ωs i)
    then 1
    else X - C (ωs i)

content: noncomputable def ElocPolyF (ωs f : Fin n → F) (p : Polynomial F) : Polynomial F :=
  ElocPoly n (liftF ωs) (liftF f) p

-- Local lemmas used:
content: @[simp]
protected lemma elocPoly_zero : ElocPoly 0 ωs f p = 1

content: @[simp]
protected lemma elocPoly_succ :
  ElocPoly (n + 1) ωs f p =
  ElocPoly n ωs f p *
    if f n = p.eval (ωs n)
    then 1
    else X - C (ωs n)

content: protected lemma elocPoly_congr {ωs' f' : ℕ → F}
  (h₁ : ∀ {m}, m < n → ωs m = ωs' m) (h₂ : ∀ {m}, m < n → f m = f' m) :
  ElocPoly n ωs f = ElocPoly n ωs' f'

content: @[simp]
protected lemma elocPolyF_eq_elocPoly' {ωs f : Fin n → F} :
  ElocPolyF ωs f p = ElocPoly n (liftF ωs) (liftF f) p

content: protected lemma elocPoly_leftF_leftF_eq_contract {ωs f : Fin m → F} :
  ElocPoly n (liftF ωs) (liftF f) =
  ElocPoly n (contract n ωs) (contract n f)

-- Target theorem and proof:
@[simp]
lemma elocPolyF_deg {ωs f : Fin n → F} : (ElocPolyF ωs f p).natDegree = Δ₀(f, p.eval ∘ ωs)
:= by
  rw [elocPolyF_eq_elocPoly']
  induction' n with n ih
  · simp only [elocPoly_zero, natDegree_one, hamming_zero_eq_dist]
    exact funext_iff.2 (Fin.elim0 ·)
  · rw [
      elocPoly_succ,
      natDegree_mul (by simp)
                    (by aesop (erase simp liftF_succ)
                              (add simp [sub_eq_zero])
                              (add safe forward (X_ne_C (liftF ωs n)))),
      elocPoly_leftF_leftF_eq_contract
    ]
    aesop (config := {warnOnNonterminal := false}) (add simp [
      hammingDist.eq_def, Finset.card_filter, Finset.sum_fin_eq_sum_range, Finset.sum_range_succ
    ]) <;> (apply Finset.sum_congr rfl; aesop (add safe (by omega)))