import Mathlib.RingTheory.Polynomial.Basic
import Goppadecoding.logic
import Goppadecoding.vanishing
import Goppadecoding.approximant
import Goppadecoding.hamming

open Polynomial

namespace Goppadecoding

section interpolation_with_errors

variable {k: Type _} [Field k] [DecidableEq k]
variable {U: Type _} [DecidableEq U]

theorem interpolation_with_errors
  {α e: U → k} {S: Finset U} {A B a b f: k[X]} {t: ℕ}
  (injective: Set.InjOn α S)
  (Adef: A = monic_vanishing_at α S)
  (appr: approximant A B t a b)
  (fdeg: 2*t + f.degree < A.degree)
  (edef: (∀ s, s ∈ S → e s = B.eval (α s) - f.eval (α s)))
  (wt: hamming_weight e S ≤ t)
: (a ∣ A ∧
   a ≠ 0 ∧
   f = B-b*(A/a) ∧
   2*t + (a*B-b*A).degree < S.card + a.degree ∧
   { s ∈ S | e s ≠ 0 } = { s ∈ S | a.eval (α s) = 0 }
  ) :=
by
  let E := monic_vanishing_at α (S.filter (fun s ↦ e s = 0))
  let c := monic_vanishing_at α (S.filter (fun s ↦ e s ≠ 0))
  have cdeg: degree c ≤ t := by calc
    _ = ((S.filter (fun s ↦ e s ≠ 0)).card: WithBot ℕ) := by apply degree_monic_vanishing_at
    _ = ((hamming_weight e S): WithBot ℕ) := by rw[hamming_weight]
    _ ≤ (t: WithBot ℕ) := by norm_cast
  have Ec: E*c = A := by rw[monic_vanishing_at_filter,Adef]
  have Ediv: E ∣ B-f := by
    apply monic_vanishing_at_divides
    . exact injon_filter α S (fun s => e s = 0) injective
    intro s es0
    have sS: s ∈ S := by apply Finset.mem_of_mem_filter s es0
    calc
      (B-f).eval (α s) = B.eval (α s) - f.eval (α s) := eval_sub B f (α s)
      _ = e s := by rw[← edef s sS]
      _ = 0 := by
            simp only [Finset.mem_filter] at es0
            simp only [es0]
  cases Ediv with | intro d BfEd =>
    have cBdA: c*B-d*A = c*f := by
      calc
        _ = c*B-d*(E*c) := by rw[← Ec]
        _ = c*B-(E*d)*c := by ring
        _ = c*B-(B-f)*c := by rw[← BfEd]
        _ = c*f := by ring
    have degcBdA: t + degree (c*B-d*A) < degree A := by
      have cdeg2: (t + degree f) + degree c ≤ (t + degree f) + t := by apply add_le_add_left cdeg
      calc
        _ = t + degree (c*f) := by rw[cBdA]
        _ = t + (degree c + degree f) := by rw[degree_mul]
        _ = (t + degree f) + degree c := by group
        _ ≤ (t + degree f) + t := cdeg2 
        _ = 2*t + degree f := by group
        _ < degree A := fdeg
    have Lexists: ∃ L, (c = a*L ∧ d = b*L) := approximant_best appr cdeg degcBdA
    cases Lexists with | intro L caLdbL =>
      have cnonzero: c ≠ 0 := by apply monic_vanishing_at_nonzero
      have ac: a ∣ c := by
        apply Dvd.intro L
        rw[caLdbL.1]
      have aA: a ∣ A := by 
        have cA: c ∣ A := by
          apply Dvd.intro E
          rw[← Ec]
          ring
        exact dvd_trans ac cA
      constructor
      . exact aA
      have anonzero: a ≠ 0 := nonzero_if_divides_nonzero ac cnonzero
      constructor
      . exact anonzero
      have afaBbA: a*f = a*B-b*A := by
        have Leq: L*(a*f) = L*(a*B-b*A) := by calc
          _ = (a*L)*f := by ring
          _ = c*f := by rw[← caLdbL.1]
          _ = c*B-d*A := by rw[← cBdA]
          _ = (a*L)*B-d*A := by rw[caLdbL.1]
          _ = (a*L)*B-(b*L)*A := by rw[caLdbL.2]
          _ = L*(a*B-b*A) := by ring
        have Ldivc: L ∣ c := by
          apply Dvd.intro a
          rw[mul_comm]
          rw[caLdbL.1]
        have Lnonzero: L ≠ 0 := by exact nonzero_if_divides_nonzero Ldivc cnonzero
        apply mul_left_cancel₀ Lnonzero Leq
      have aAa: a*(A/a) = A := EuclideanDomain.mul_div_cancel' anonzero aA
      have fBbAa: f = B-b*(A/a) := by
        have fBbAa1: a*f = a*(B-b*(A/a)) := by calc
          a*f = a*B-b*A := afaBbA
          _ = a*B-b*(a*(A/a)) := by rw[aAa]
          _ = a*(B-b*(A/a)) := by ring
        apply mul_left_cancel₀ anonzero fBbAa1
      constructor
      . exact fBbAa
      constructor
      . calc
          _ = 2*t + degree (a*f) := by rw[afaBbA]
          _ = 2*t + (degree a + degree f) := by rw[degree_mul]
          _ = (2*t + degree f) + degree a := by group
          _ < degree A + degree a := by apply add_lt_add_right_nonzero_degree anonzero fdeg
          _ = degree (monic_vanishing_at α S) + degree a := by rw[Adef]
          _ = S.card + degree a := by rw[degree_monic_vanishing_at]
      ext s
      simp only [ne_eq, Set.mem_setOf_eq, and_congr_right_iff]
      intro sinS
      rw[edef s sinS]
      have Aeval: A.eval (α s) = 0 := by
        rw[Adef]
        apply monic_vanishing_at_vanishes α S sinS
      by_cases aeval: a.eval (α s) = 0
      . simp only [aeval, iff_true, ne_eq]
        have Aaeval: (A/a).eval (α s) ≠ 0 := by
          have derivA: a*(A/a).diff + a.diff*(A/a) = A.diff := by calc
            _ = (a*(A/a)).diff := by
                  rw[diff,add_comm]
                  simp only [AddHom.toFun_eq_coe, LinearMap.coe_toAddHom, derivative_mul]
            _ = A.diff := by rw[aAa]
          have derivAeval: a.diff.eval (α s) * (A/a).eval (α s) ≠ 0 := by calc
            _ = 0 * (A/a).diff.eval (α s) + a.diff.eval (α s) * (A/a).eval (α s) := by ring
            _ = a.eval (α s) * (A/a).diff.eval (α s) + a.diff.eval (α s) * (A/a).eval (α s) := by simp only [zero_mul, zero_add, aeval]
            _ = (a*(A/a).diff).eval (α s) + (a.diff*(A/a)).eval (α s) := by repeat rw[eval_mul]
            _ = (a*(A/a).diff + a.diff*(A/a)).eval (α s) := by rw[eval_add]
            _ = A.diff.eval (α s) := by rw[derivA]
            _ = (monic_vanishing_at α S).diff.eval (α s) := by rw[Adef]
            _ ≠ 0 := by apply derivative_monic_vanishing_at_eval_nonzero α injective sinS
          exact right_ne_zero_of_mul derivAeval
        have beval: b.eval (α s) ≠ 0 := by
          let coprime := appr.1
          exact no_shared_root_if_coprime coprime aeval
        have bAaeval: b.eval (α s) * (A/a).eval (α s) ≠ 0 := mul_ne_zero beval Aaeval
        have Bfeval: B.eval (α s) - f.eval (α s) ≠ 0 := by calc
          _ = (B-f).eval (α s) := by rw[eval_sub]
          _ = (B-(B-b*(A/a))).eval (α s) := by rw[fBbAa]
          _ = (b*(A/a)).eval (α s) := by group
          _ = b.eval (α s) * (A/a).eval (α s) := by rw[eval_mul]
          _ ≠ 0 := bAaeval
        simp only [Bfeval, not_false_eq_true]
      . simp only [aeval, iff_false, not_not]
        have aBfeval: a.eval (α s) * (B.eval (α s) - f.eval (α s)) = 0 := by calc
          _ = (a*(B-f)).eval (α s) := by rw[← eval_sub,← eval_mul]
          _ = (a*B-a*f).eval (α s) := by ring_nf
          _ = (a*B-(a*B-b*A)).eval (α s) := by rw[afaBbA]
          _ = (b*A).eval (α s) := by ring_nf
          _ = b.eval (α s) * A.eval (α s) := by rw[eval_mul]
          _ = b.eval (α s) * 0 := by rw[Aeval]
          _ = 0 := by ring
        rw[← not_eq_iff] at aeval
        apply zero_if_mul_left_zero aeval aBfeval

theorem checking_interpolation_with_errors
  {α: U → k} {S: Finset U} {A B a b f: k[X]} {t: ℕ}
  (injective: Set.InjOn α S)
  (Adef: A = monic_vanishing_at α S)
  (adeg: a.degree ≤ t)
  (aA: a ∣ A)
  (aBbAdeg: 2*t + (a*B-b*A).degree < S.card + a.degree)
  (fdef: f = B-b*(A/a))
: (a*(B-f) = b*A ∧
   2*t + f.degree < S.card ∧
   hamming_weight (fun s ↦ B.eval (α s) - f.eval (α s)) S ≤ t
  ) :=
by
  have Anonzero: A ≠ 0 := by rw[Adef]; exact monic_vanishing_at_nonzero α S
  have anonzero: a ≠ 0 := nonzero_if_divides_nonzero aA Anonzero
  have aAa: a*(A/a) = A := EuclideanDomain.mul_div_cancel' anonzero aA
  have af: a*f = a*B-b*A := by calc
    a*f = a*(B-b*(A/a)) := by rw[fdef]
    _ = a*B-b*(a*(A/a)) := by ring
    _ = a*B-b*A := by rw[aAa]
  constructor
  . calc
      a*(B-f) = a*B-a*f := by ring
      _ = a*B-(a*B-b*A) := by rw[af]
      _ = b*A := by ring
  constructor
  . have afdeg: (2*t + degree f) + degree a < S.card + degree a := by calc
      _ = 2*t + (degree a + degree f) := by group
      _ = 2*t + degree (a*f) := by rw[← degree_mul]
      _ = 2*t + degree (a*B-b*A) := by rw[af]
      _ < S.card + degree a := aBbAdeg
    exact lt_of_add_lt_add_right afdeg
  unfold hamming_weight
  simp only [ne_eq]
  have Bfroots: (S.filter (fun s ↦ ¬B.eval (α s) - f.eval (α s) = 0)) ⊆ (S.filter (fun s ↦ a.eval (α s) = 0)) := by
    intro s
    simp only [Finset.mem_filter, and_imp]
    intro sinS Bfnonzero
    rw[← not_eq_iff] at Bfnonzero
    constructor
    . exact sinS
    have aBfeval: a.eval (α s) * (B.eval (α s) - f.eval (α s)) = 0 := by calc
      _ = (a*(B-f)).eval (α s) := by rw[eval_mul,eval_sub]
      _ = (a*(B-(B-b*(A/a)))).eval (α s) := by rw[fdef]
      _ = (b*(a*(A/a))).eval (α s) := by ring_nf
      _ = (b*A).eval (α s) := by rw[aAa]
      _ = b.eval (α s) * A.eval (α s) := by rw[eval_mul]
      _ = b.eval (α s) * (monic_vanishing_at α S).eval (α s) := by rw[Adef]
      _ = b.eval (α s) * 0 := by rw[monic_vanishing_at_vanishes α S sinS]
      _ = 0 := by ring
    exact zero_if_mul_right_zero Bfnonzero aBfeval
  have Bfroots2: (S.filter (fun s ↦ ¬B.eval (α s) - f.eval (α s) = 0)).card ≤ (S.filter (fun s ↦ a.eval (α s) = 0)).card := Finset.card_le_of_subset Bfroots
  have rootst: (S.filter (fun s ↦ a.eval (α s) = 0)).card ≤ (t:WithBot ℕ) := by calc
    _ ≤ degree a := card_filter_image_roots_le_degree α S a injective anonzero
    _ ≤ t := adeg
  have rootst2: (S.filter (fun s ↦ a.eval (α s) = 0)).card ≤ t := WithBot.coe_le_coe.mp rootst
  exact le_trans Bfroots2 rootst2

end interpolation_with_errors