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

namespace UniPoly

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

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

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

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

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

-- 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: lemma coeff_eq_zero {p : UniPoly Q} :
    (∀ i, (hi : i < p.size) → p[i] = 0) ↔ ∀ i, p.coeff i = 0

-- Target theorem and proof:
lemma eq_degree_of_equiv [LawfulBEq R] {p q : UniPoly R} : equiv p q → p.degree = q.degree
:= by
  unfold equiv degree
  intro h_equiv
  induction p using last_nonzero_induct with
  | case1 p h_none_p h_all_zero =>
    have h_zero_p : ∀ i, p.coeff i = 0 := coeff_eq_zero.mp h_all_zero
    have h_zero_q : ∀ i, q.coeff i = 0 := by intro i; rw [← h_equiv, h_zero_p]
    have h_none_q : q.last_nonzero = none := last_nonzero_none (coeff_eq_zero.mpr h_zero_q)
    rw [h_none_p, h_none_q]
  | case2 p k h_some_p h_nonzero_p h_max_p =>
    have h_equiv_k := h_equiv k
    have k_lt_q : k < q.size := by
      by_contra h_not_lt
      have h_ge := Nat.le_of_not_lt h_not_lt
      simp [h_ge] at h_equiv_k
      contradiction
    simp [k_lt_q] at h_equiv_k
    have h_nonzero_q : q[k.val] ≠ 0 := by rwa [← h_equiv_k]
    have h_max_q : ∀ j, (hj : j < q.size) → j > k → q[j] = 0 := by
      intro j hj j_gt_k
      have h_eq := h_equiv j
      simp [hj] at h_eq
      rw [← h_eq]
      rcases Nat.lt_or_ge j p.size with hj | hj
      · simp [hj, h_max_p j hj j_gt_k]
      · simp [hj]
    have h_some_q : q.last_nonzero = some ⟨ k, k_lt_q ⟩ :=
      last_nonzero_some_iff.mpr ⟨ h_nonzero_q, h_max_q ⟩
    rw [h_some_p, h_some_q]