-- ArkLib/Data/UniPoly/Basic.lean

import Mathlib.Algebra.Tropical.Basic

import Mathlib.RingTheory.Polynomial.Basic

import ArkLib.Data.Array.Lemmas

open Polynomial

@[reducible, inline, specialize]
def UniPoly (R : Type*) := Array R

def Polynomial.toImpl {R : Type*} [Semiring R] (p : R[X]) : UniPoly R :=
  match p.degree with
  | ⊥ => #[]
  | some d  => .ofFn (fun i : Fin (d + 1) => p.coeff i)

namespace UniPoly

@[reducible]
def mk {R : Type*} (coeffs : Array R) : UniPoly R := coeffs

variable {R : Type*} [Ring R] [BEq R]

variable {Q : Type*} [Ring Q]

@[reducible]
def coeff (p : UniPoly Q) (i : ℕ) : Q := p.getD i 0

def last_nonzero (p : UniPoly R) : Option (Fin p.size) :=
  p.findIdxRev? (· != 0)

def trim (p : UniPoly R) : UniPoly R :=
  match p.last_nonzero with
  | none => #[]
  | some i => p.extract 0 (i.val + 1)

def degree (p : UniPoly R) : Nat :=
  match p.last_nonzero with
  | none => 0
  | some i => i.val + 1

namespace Trim

def last_nonzero_prop {p : UniPoly R} (k : Fin p.size) : Prop :=
  p[k] ≠ 0 ∧ (∀ j, (hj : j < p.size) → j > k → p[j] = 0)

def equiv (p q : UniPoly R) : Prop :=
  ∀ i, p.coeff i = q.coeff i

end Trim

def UniPolyC (R : Type*) [BEq R] [Ring R] := { p : UniPoly R // p.trim = p }

section Operations

variable {S : Type*}

def eval₂ [Semiring S] (f : R →+* S) (x : S) (p : UniPoly R) : S :=
  p.zipIdx.foldl (fun acc ⟨a, i⟩ => acc + f a * x ^ i) 0

@[inline, specialize]
def add_raw (p q : UniPoly R) : UniPoly R :=
  let ⟨p', q'⟩ := Array.matchSize p q 0
  .mk (Array.zipWith (· + ·) p' q' )

variable (p q r : UniPoly R)

def canonical (p : UniPoly R) := p.trim = p

end Operations

namespace OperationsC

variable {R : Type*} [Ring R] [BEq R] [LawfulBEq R]

variable (p q r : UniPolyC R)

end OperationsC

section ToPoly

noncomputable def toPoly (p : UniPoly R) : Polynomial R :=
  p.eval₂ Polynomial.C Polynomial.X

noncomputable def UniPolyC.toPoly (p : UniPolyC R) : Polynomial R := p.val.toPoly

alias ofPoly := Polynomial.toImpl

-- Repository definitions used:
content: def findIdxRev? (cond : α → Bool) (as : Array α) : Option (Fin as.size) :=
  find ⟨ as.size, Nat.lt_succ_self _ ⟩
where
  find : Fin (as.size + 1) → Option (Fin as.size)
    | 0 => none
    | ⟨ i+1, h ⟩ =>
      if (cond as[i]) then
        some ⟨ i, Nat.lt_of_succ_lt_succ h ⟩
      else
        find ⟨ i, Nat.lt_of_succ_lt h ⟩

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

content: @[reducible]
def matchSize (a b : Array α) (unit : α) : Array α × Array α :=
  (a.rightpad (b.size) unit, b.rightpad (a.size) unit)

content: def findIdxRev?_maximal {cond} {as : Array α} {k : Fin as.size} :
  findIdxRev? cond as = some k → ∀ j : Fin as.size, j > k → ¬ cond as[j] :=

content: def findIdxRev?_def {cond} {as : Array α} {k : Fin as.size} :
  findIdxRev? cond as = some k → cond as[k] :=

content: def getLast (a : Array α) (h : a.size > 0) : α := a[a.size - 1]

-- Repository lemmas used:
content: theorem findIdxRev?_eq_some {cond} {as : Array α} (h : ∃ i, ∃ hi : i < as.size, cond as[i]) :
  ∃ k : Fin as.size, findIdxRev? cond as = some k

content: theorem findIdxRev?_eq_none {cond} {as : Array α} (h : ∀ i, (hi : i < as.size) → ¬ cond as[i]) :
  findIdxRev? cond as = none

content: theorem findIdxRev?_emtpy_none {cond} {as : Array α} (h : as = #[]) :
  findIdxRev? cond as = none

-- Local definitions used:
content: @[reducible, inline, specialize]
def UniPoly (R : Type*) := Array R

content: def Polynomial.toImpl {R : Type*} [Semiring R] (p : R[X]) : UniPoly R :=
  match p.degree with
  | ⊥ => #[]
  | some d  => .ofFn (fun i : Fin (d + 1) => p.coeff i)

content: @[reducible]
def mk {R : Type*} (coeffs : Array R) : UniPoly R := coeffs

content: @[reducible]
def coeff (p : UniPoly Q) (i : ℕ) : Q := p.getD i 0

content: def last_nonzero (p : UniPoly R) : Option (Fin p.size) :=
  p.findIdxRev? (· != 0)

content: def trim (p : UniPoly R) : UniPoly R :=
  match p.last_nonzero with
  | none => #[]
  | some i => p.extract 0 (i.val + 1)

content: def degree (p : UniPoly R) : Nat :=
  match p.last_nonzero with
  | none => 0
  | some i => i.val + 1

content: def last_nonzero_prop {p : UniPoly R} (k : Fin p.size) : Prop :=
  p[k] ≠ 0 ∧ (∀ j, (hj : j < p.size) → j > k → p[j] = 0)

content: def equiv (p q : UniPoly R) : Prop :=
  ∀ i, p.coeff i = q.coeff i

content: def UniPolyC (R : Type*) [BEq R] [Ring R] := { p : UniPoly R // p.trim = p }

content: def eval₂ [Semiring S] (f : R →+* S) (x : S) (p : UniPoly R) : S :=
  p.zipIdx.foldl (fun acc ⟨a, i⟩ => acc + f a * x ^ i) 0

content: @[inline, specialize]
def add_raw (p q : UniPoly R) : UniPoly R :=
  let ⟨p', q'⟩ := Array.matchSize p q 0
  .mk (Array.zipWith (· + ·) p' q' )

content: def canonical (p : UniPoly R) := p.trim = p

content: noncomputable def toPoly (p : UniPoly R) : Polynomial R :=
  p.eval₂ Polynomial.C Polynomial.X

content: noncomputable def UniPolyC.toPoly (p : UniPolyC R) : Polynomial R := p.val.toPoly

alias ofPoly := Polynomial.toImpl

-- Local lemmas used:
content: theorem last_nonzero_none [LawfulBEq R] {p : UniPoly R} :
  (∀ i, (hi : i < p.size) → p[i] = 0) → p.last_nonzero = none

content: theorem last_nonzero_some [LawfulBEq R] {p : UniPoly R} {i} (hi : i < p.size) (h : p[i] ≠ 0) :
  ∃ k, p.last_nonzero = some k

content: theorem last_nonzero_spec [LawfulBEq R] {p : UniPoly R} {k} :
  p.last_nonzero = some k
  → p[k] ≠ 0 ∧ (∀ j, (hj : j < p.size) → j > k → p[j] = 0)

content: lemma last_nonzero_unique {p : UniPoly Q} {k k' : Fin p.size} :
  last_nonzero_prop k → last_nonzero_prop k' → k = k'

content: theorem last_nonzero_some_iff [LawfulBEq R] {p : UniPoly R} {k} :
  p.last_nonzero = some k ↔ (p[k] ≠ 0 ∧ (∀ j, (hj : j < p.size) → j > k → p[j] = 0))

content: theorem last_nonzero_induct [LawfulBEq R] {motive : UniPoly R → Prop}
  (case1 : ∀ p, p.last_nonzero = none → (∀ i, (hi : i < p.size) → p[i] = 0) → motive p)
  (case2 : ∀ p : UniPoly R, ∀ k : Fin p.size, p.last_nonzero = some k → p[k] ≠ 0 →
    (∀ j : ℕ, (hj : j < p.size) → j > k → p[j] = 0) → motive p)
  (p : UniPoly R) : motive p

content: theorem induct [LawfulBEq R] {motive : UniPoly R → Prop}
  (case1 : ∀ p, p.trim = #[] → (∀ i, (hi : i < p.size) → p[i] = 0) → motive p)
  (case2 : ∀ p : UniPoly R, ∀ k : Fin p.size, p.trim = p.extract 0 (k + 1)
    → p[k] ≠ 0 → (∀ j : ℕ, (hj : j < p.size) → j > k → p[j] = 0) → motive p)
  (p : UniPoly R) : motive p

content: theorem size_eq_degree (p : UniPoly R) : p.trim.size = p.degree

content: theorem size_le_size (p : UniPoly R) : p.trim.size ≤ p.size

content: theorem coeff_eq_getElem_of_lt [LawfulBEq R] {p : UniPoly R} {i} (hi : i < p.size) :
  p.trim.coeff i = p[i]

content: theorem coeff_eq_coeff [LawfulBEq R] (p : UniPoly R) (i : ℕ) :
  p.trim.coeff i = p.coeff i

content: lemma coeff_eq_getElem {p : UniPoly Q} {i} (hp : i < p.size) :
  p.coeff i = p[i]

content: lemma coeff_eq_zero {p : UniPoly Q} :
    (∀ i, (hi : i < p.size) → p[i] = 0) ↔ ∀ i, p.coeff i = 0

content: lemma eq_degree_of_equiv [LawfulBEq R] {p q : UniPoly R} : equiv p q → p.degree = q.degree

content: theorem eq_of_equiv [LawfulBEq R] {p q : UniPoly R} : equiv p q → p.trim = q.trim

content: theorem canonical_empty : (UniPoly.mk (R:=R) #[]).trim = #[]

content: theorem canonical_of_size_zero {p : UniPoly R} : p.size = 0 → p.trim = p

content: theorem canonical_nonempty_iff [LawfulBEq R] {p : UniPoly R} (hp : p.size > 0) :
  p.trim = p ↔ p.last_nonzero = some ⟨ p.size - 1, Nat.pred_lt_self hp ⟩

content: theorem last_nonzero_last_iff [LawfulBEq R] {p : UniPoly R} (hp : p.size > 0) :
  p.last_nonzero = some ⟨ p.size - 1, Nat.pred_lt_self hp ⟩ ↔ p.getLast hp ≠ 0

content: @[ext] theorem UniPolyC.ext {p q : UniPolyC R} (h : p.val = q.val) : p = q

content: lemma matchSize_size_eq {p q : UniPoly Q} :
    let (p', q') := Array.matchSize p q 0
    p'.size = q'.size

content: lemma matchSize_size {p q : UniPoly Q} :
    let (p', _) := Array.matchSize p q 0
    p'.size = max p.size q.size

content: lemma zipWith_size {R} {f : R → R → R} {a b : Array R} (h : a.size = b.size) :
    (Array.zipWith f a b).size = a.size

content: theorem add_size {p q : UniPoly Q} : (add_raw p q).size = max p.size q.size

content: theorem add_coeff {p q : UniPoly Q} {i : ℕ} (hi : i < (add_raw p q).size) :
  (add_raw p q)[i] = p.coeff i + q.coeff i

content: theorem zero_add (hp : p.canonical) : 0 + p = p

content: lemma coeff_toPoly {p : UniPoly Q} {n : ℕ} : p.toPoly.coeff n = p.coeff n

content: lemma toImpl_elim (p : Q[X]) :
    (p = 0 ∧ p.toImpl = #[])
  ∨ (p ≠ 0 ∧ p.toImpl = .ofFn (fun i : Fin (p.natDegree + 1) => p.coeff i))

content: theorem toPoly_toImpl {p : Q[X]} : p.toImpl.toPoly = p

content: theorem trim_toImpl [LawfulBEq R] (p : R[X]) : p.toImpl.trim = p.toImpl

-- Target theorem and proof:
lemma toImpl_toPoly_of_canonical [LawfulBEq R] (p : UniPolyC R) : p.toPoly.toImpl = p
:= by
  -- we will change something slightly more general: `toPoly` is injective on canonical polynomials
  suffices h_inj : ∀ q : UniPolyC R, p.toPoly = q.toPoly → p = q by
    have : p.toPoly = p.toPoly.toImpl.toPoly := by rw [toPoly_toImpl]
    exact h_inj ⟨ p.toPoly.toImpl, trim_toImpl p.toPoly ⟩ this |> congrArg Subtype.val |>.symm
  intro q hpq
  apply UniPolyC.ext
  apply Trim.canonical_ext p.property q.property
  intro i
  rw [← coeff_toPoly, ← coeff_toPoly]
  exact hpq |> congrArg (fun p => p.coeff i)