import Mathlib.RingTheory.Polynomial.Basic
import Goppadecoding.finset
import Goppadecoding.dvd
import Goppadecoding.poly
import Goppadecoding.diff

namespace Goppadecoding

open Polynomial BigOperators Finset

section vanishing

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

noncomputable def monic_vanishing_at (α: U → k) (S: Finset U): k[X] := 
  ∏ s in S, (X - C (α s))

theorem monic_vanishing_at_filter (α: U → k) (S: Finset U) (p: U → Prop) [DecidablePred p]:
  monic_vanishing_at α (S.filter p) * monic_vanishing_at α (S.filter (fun s ↦ ¬(p s))) = monic_vanishing_at α S := by
  unfold monic_vanishing_at
  apply prod_filter_mul_prod_filter_not

-- meaningful only for s in S 
noncomputable def monic_vanishing_at_except (α: U → k) (S: Finset U) (s: U): k[X] := 
  monic_vanishing_at α (S.erase s)

-- meaningful only for s in S
noncomputable def vanishing_at_except_1_at (α: U → k) (S: Finset U) (s: U): k[X] :=
  (monic_vanishing_at_except α S s) * C (((monic_vanishing_at_except α S s).eval (α s))⁻¹)

noncomputable def interpolator (α: U → k) (S: Finset U) (r: U → k): k[X] :=
  ∑ s in S, C (r s) * vanishing_at_except_1_at α S s

theorem monic_vanishing_at_eval (α: U → k) (S: Finset U) (z: k):
  (monic_vanishing_at α S).eval (z) = ∏ s in S, (z - α s) := by
  unfold monic_vanishing_at
  simp only [eval_prod, eval_sub, eval_X, eval_C]

theorem monic_vanishing_at_vanishes (α: U → k) (S: Finset U) {s: U}:
  s ∈ S → (monic_vanishing_at α S).eval (α s) = 0 := by
  rw[monic_vanishing_at_eval]
  intro sinS
  apply prod_eq_zero sinS
  simp only [sub_self]

theorem vanishing_at_except_1_at_vanishes (α: U → k) (S: Finset U) {s: U} {t: U}:
  t ∈ S → t ≠ s → (vanishing_at_except_1_at α S s).eval (α t) = 0 := by
  unfold vanishing_at_except_1_at
  intro tinS distinct
  have tinSerases: t ∈ S.erase s := mem_erase_of_ne_of_mem distinct tinS
  rw[eval_mul]
  unfold monic_vanishing_at_except
  rw[monic_vanishing_at_vanishes α (S.erase s) tinSerases]
  ring
  
theorem monic_vanishing_at_except_eval_nonzero (α: U → k) (S: Finset U) (s: U):
  Set.InjOn α S → s ∈ S → (monic_vanishing_at_except α S s).eval (α s) ≠ 0 := by
  unfold monic_vanishing_at_except
  intro injective sinS
  rw[monic_vanishing_at_eval]
  simp only [ne_eq, prod_eq_zero_iff, mem_erase, not_exists, not_and, and_imp]
  intro t tnes tinS
  have distinctij: α t ≠ α s := by apply (Set.InjOn.ne_iff injective tinS sinS).mpr tnes
  have nonzerodiff: α s - α t ≠ 0 := sub_ne_zero.mpr (Ne.symm distinctij)
  simp only [nonzerodiff, not_false_eq_true]

theorem vanishing_at_except_1_at_eval_1 (α: U → k) (S: Finset U) (s: U):
  Set.InjOn α S → s ∈ S → (vanishing_at_except_1_at α S s).eval (α s) = 1 := by
  unfold vanishing_at_except_1_at
  rw[eval_mul,eval_C]
  intro injective sinS
  apply mul_inv_cancel
  apply monic_vanishing_at_except_eval_nonzero
  apply injective 
  apply sinS
  
theorem vanishing_at_except_1_at_eval {α: U → k} {S: Finset U} {s: U} {t: U}:
  Set.InjOn α S → s ∈ S → t ∈ S → 
  (vanishing_at_except_1_at α S s).eval (α t) = if t = s then 1 else 0 := by
  intro injective sinS tinS
  by_cases tvss: t = s
  . simp only [tvss, injective, sinS, vanishing_at_except_1_at_eval_1, ite_true]
  . simp only [tinS, ne_eq, tvss, not_false_eq_true, vanishing_at_except_1_at_vanishes, ite_false]

theorem degree_monic_vanishing_at (α: U → k) (S: Finset U):
  (monic_vanishing_at α S).degree = S.card := by
  unfold monic_vanishing_at
  simp only [degree_prod, degree_X_sub_C, Finset.sum_const, nsmul_eq_mul, mul_one]

theorem degree_monic_vanishing_at_except {α: U → k} {S: Finset U} {s: U}:
  s ∈ S → (monic_vanishing_at_except α S s).degree < S.card := by
  unfold monic_vanishing_at_except
  rw[degree_monic_vanishing_at]
  norm_cast
  intro sinS
  exact card_erase_lt_of_mem sinS

theorem monic_vanishing_at_nonzero (α: U → k) (S: Finset U):
  monic_vanishing_at α S ≠ 0 := by
  by_cases z: monic_vanishing_at α S = 0
  . have cardbot: (S.card: WithBot ℕ) = ⊥ := by calc
      S.card = (monic_vanishing_at α S).degree := by rw[← degree_monic_vanishing_at α S]
      _ = (0:k[X]).degree := by rw[z]
      _ = ⊥ := by simp only [degree_zero]
    have cardnotbot: (S.card: WithBot ℕ) ≠ ⊥ := WithBot.nat_ne_bot (S.card)
    exact False.elim (cardnotbot cardbot)
  . simp only [ne_eq, z, not_false_eq_true]

theorem degree_vanishing_at_except_1_at (α: U → k) (S: Finset U) (s: U):
  s ∈ S → (vanishing_at_except_1_at α S s).degree < S.card := by
  unfold vanishing_at_except_1_at
  intro sinS
  apply degree_mul_C_lt_if_degree_lt
  unfold monic_vanishing_at_except
  rw[degree_monic_vanishing_at]
  rw[card_erase_of_mem sinS]
  apply Nat.cast_lt.mpr
  apply Nat.sub_lt
  have nonzerocard: S.card ≠ 0 := card_ne_zero_of_mem sinS
  exact Nat.pos_of_ne_zero nonzerocard
  norm_num

theorem degree_interpolator_lt (α: U → k) (S: Finset U) (r: U → k):
  (interpolator α S r).degree < S.card :=
by
  unfold interpolator
  apply sum_degree_lt
  intro j jinS
  apply degree_C_mul_lt_if_degree_lt
  apply degree_vanishing_at_except_1_at
  simp only[jinS]

theorem monic_vanishing_at_empty (α: U → k):
  monic_vanishing_at α ∅ = 1 := by
  unfold monic_vanishing_at
  simp only [Finset.prod_empty]

theorem monic_vanishing_at_insert (α: U → k) (S: Finset U) (u: U):
  u ∉ S → monic_vanishing_at α (insert u S) = (X - C (α u)) * (monic_vanishing_at α S) := by
  unfold monic_vanishing_at
  intro unotinS
  simp only [unotinS, not_false_eq_true, prod_insert]

theorem monic_vanishing_at_divides {α: U → k}:
  ∀ S: Finset U, ∀ f: k[X],
  Set.InjOn α S →
  ((∀ s, (s ∈ S → f.eval (α s) = 0)) → monic_vanishing_at α S ∣ f) := by
  apply Finset.induction
  . rw[monic_vanishing_at_empty]
    simp only [coe_empty, Set.injOn_empty, not_mem_empty, IsEmpty.forall_iff, implies_true, isUnit_one, IsUnit.dvd, forall_true_left, forall_const]
  . intro u S unotinS prevcase f injective roots
    rw[monic_vanishing_at_insert α S u unotinS]
    have uroot: f.eval (α u) = 0 := by simp only [mem_insert, true_or, roots]
    have XCdiv: X - C (α u) ∣ f := dvd_iff_isRoot.mpr uroot
    cases XCdiv with
    | intro q qf =>
      have qroots: ∀ s, (s ∈ S → q.eval (α s) = 0) := by
        intro s sinS
        have step1: (α s - α u) * q.eval (α s) = 0 := by
          calc
            (α s - α u) * q.eval (α s)
              = (X - C (α u)).eval (α s) * q.eval (α s) := by simp only [eval_sub, eval_X, eval_C]
            _ = ((X - C (α u)) * q).eval (α s) := by simp only [eval_sub, eval_X, eval_C, eval_mul]
            _ = f.eval (α s) := by rw[qf]
            _ = 0 := by apply roots s (mem_insert_of_mem sinS)
        by_cases su: s = u 
        . rw[su] at sinS 
          exact False.elim (unotinS sinS)
        . have uinSinsert: u ∈ insert u S := by apply mem_insert_self
          have sinSinsert: s ∈ insert u S := by apply mem_insert_of_mem sinS
          have sneu: s ≠ u := by simp only [ne_eq, su, not_false_eq_true]
          have asu: α s ≠ α u := by apply (Set.InjOn.ne_iff injective sinSinsert uinSinsert).mpr sneu
          have asuminus: α s - α u ≠ 0 := sub_ne_zero.mpr asu
          by_cases qzero: q.eval (α s) = 0
          . simp only[qzero]
          have step1false: (α s - α u) * q.eval (α s) ≠ 0 := mul_ne_zero asuminus qzero
          exact False.elim (step1false step1)
      have subsetinsert: S ⊆ insert u S := by apply Finset.subset_insert
      have previnjective: Set.InjOn α S := Set.InjOn.mono subsetinsert injective
      have dvdq: monic_vanishing_at α S ∣ q := prevcase q previnjective qroots
      rw[qf]
      apply mul_dvd_mul_left
      apply dvdq

theorem vanishing_at_degree (α: U → k) (S: Finset U) (f: k[X]):
  Set.InjOn α S → 
  (∀ s, (s ∈ S → f.eval (α s) = 0)) → (f = 0 ∨ f.degree ≥ S.card) := by
  intro injective roots
  have div: monic_vanishing_at α S ∣ f := by 
    apply monic_vanishing_at_divides S f injective
    apply roots
  by_cases fsign: f ≠ 0 
  . have degdeg: (monic_vanishing_at α S).degree ≤ f.degree := degree_le_of_dvd div fsign
    rw[degree_monic_vanishing_at] at degdeg
    simp only [ge_iff_le, degdeg, or_true]
  . left 
    apply not_ne_iff.mp
    apply fsign

theorem degree_below_vanishing (α: U → k) (S: Finset U) (f: k[X]):
  Set.InjOn α S → 
  (∀ s, (s ∈ S → f.eval (α s) = 0)) → f.degree < S.card → f = 0 := by
  intro injective roots
  have v: f = 0 ∨ f.degree ≥ S.card := by apply vanishing_at_degree α S f injective roots
  contrapose
  rename_i inst inst_1
  intro a
  simp_all only [ge_iff_le, false_or, not_lt]

-- maybe merge this with monic_vanishing_at_insert
theorem monic_vanishing_at_except_filter (α: U → k) (S: Finset U) (s: U):
  s ∈ S → monic_vanishing_at_except α S s * (X - C (α s)) = monic_vanishing_at α S := by
  unfold monic_vanishing_at_except
  unfold monic_vanishing_at
  apply prod_erase_mul

theorem derivative_monic_vanishing_at (α: U → k) {S: Finset U}:
  (monic_vanishing_at α S).diff
  = ∑ s in S, monic_vanishing_at_except α S s
:= by
  unfold monic_vanishing_at_except
  unfold monic_vanishing_at
  unfold Finset.prod
  unfold Finset.sum
  simp only [diff, AddHom.toFun_eq_coe, LinearMap.coe_toAddHom, erase_val, mem_val]
  rw[derivative_prod]
  -- it's weird that a single simp doesn't finish this
  simp only [map_sub, derivative_X, derivative_C, sub_zero, mul_one, erase_val, mem_val]
  apply Finset.sum_congr
  simp only
  intro s _
  congr
  simp only [eq_iff_true_of_subsingleton]

theorem derivative_monic_vanishing_at_eval {α: U → k} {S: Finset U} {s: U}:
  s ∈ S → (monic_vanishing_at α S).diff.eval (α s) = (monic_vanishing_at_except α S s).eval (α s) := by
  intro sinS
  let A:k[X] := monic_vanishing_at_except α S s
  let Z:k[X] := X - C (α s)
  have zeval: Z.eval (α s) = 0 := by calc
    _ = X.eval (α s) - (C (α s)).eval (α s) := by rw[eval_sub]
    _ = α s - (C (α s)).eval (α s) := by rw[eval_X]
    _ = α s - α s := by rw[eval_C]
    _ = 0 := by group
  calc
    _ = (A * Z).diff.eval (α s) := by rw[monic_vanishing_at_except_filter α S s sinS]
    _ = (A.diff * Z + A * Z.diff).eval (α s) := by simp only [diff, AddHom.toFun_eq_coe, LinearMap.coe_toAddHom, derivative_mul, map_sub, derivative_X, derivative_C, sub_zero, mul_one, eval_add, eval_mul, eval_sub, eval_X, eval_C, sub_self, mul_zero, zero_add]
    _ = (A.diff * Z).eval (α s) + (A * Z.diff).eval (α s) := by rw[eval_add]
    _ = A.diff.eval (α s) * Z.eval (α s) + A.eval (α s) * Z.diff.eval (α s) := by repeat rw[eval_mul]
    _ = A.diff.eval (α s) * 0 + A.eval (α s) * Z.diff.eval (α s) := by rw[zeval]
    _ = A.eval (α s) * Z.diff.eval (α s) := by ring
    _ = A.eval (α s) * (1:k[X]).eval (α s) := by rw[diff]; simp only [AddHom.toFun_eq_coe, LinearMap.coe_toAddHom, map_sub, derivative_X, derivative_C, sub_zero, eval_one, mul_one]
    _ = A.eval (α s) * (C 1:k[X]).eval (α s) := rfl
    _ = A.eval (α s) * (1:k) := by rw[eval_C]
    _ = A.eval (α s) := by ring

theorem derivative_monic_vanishing_at_eval_nonzero (α: U → k) {S: Finset U} {s: U}:
  Set.InjOn α S → s ∈ S → (monic_vanishing_at α S).diff.eval (α s) ≠ 0 := by
  intro injective sinS
  calc
    _ = (monic_vanishing_at_except α S s).eval (α s) := by rw[derivative_monic_vanishing_at_eval sinS]
    _ ≠ 0 := by apply monic_vanishing_at_except_eval_nonzero α S s injective sinS

-- can also obtain from Polynomial.card_roots, card_image_of_injOn, etc.
theorem card_filter_image_roots_le_degree (α: U → k) [DecidableEq k] (S: Finset U) (f: k[X]):
  Set.InjOn α S → f ≠ 0 → (S.filter (fun s ↦ f.eval (α s) = 0)).card ≤ f.degree := by 
  intro injective fnonzero
  let R := S.filter (fun s ↦ f.eval (α s) = 0)
  have cardR: (monic_vanishing_at α R).degree = R.card := degree_monic_vanishing_at α R
  let injR: Set.InjOn α R := injon_filter α S (fun s ↦ f.eval (α s) = 0) injective
  let rootsR: (∀ r, r ∈ R → f.eval (α r) = 0) := by
    intro r rinR
    rw[mem_filter] at rinR
    simp only[rinR]
  have divides: monic_vanishing_at α R ∣ f := monic_vanishing_at_divides R f injR rootsR
  calc
    _ = (R.card: WithBot ℕ) := rfl
    _ = (monic_vanishing_at α R).degree := by rw[cardR]
    _ ≤ f.degree := degree_le_of_dvd divides fnonzero

theorem coprime_monic_vanishing_at_nonroots {α: U → k} [DecidableEq k] {S: Finset U} {a: k[X]}:
  IsCoprime a (monic_vanishing_at α (S.filter (fun s ↦ ¬a.eval (α s) = 0))) := by
  unfold monic_vanishing_at
  rw[IsCoprime.prod_right_iff]
  intro s
  simp only [mem_filter, and_imp]
  intro _ nonroot
  apply coprime_X_minus_nonroot
  simp only [ne_eq, nonroot, not_false_eq_true]
  
theorem leadCoeff_monic_vanishing_at {α: U → k} {S: Finset U}:
  (monic_vanishing_at α S).leadingCoeff = 1 := by
  unfold monic_vanishing_at
  simp only[leadingCoeff_prod,leadingCoeff_X_sub_C,prod_const_one]

theorem lead_mul_monic_vanishing_at_roots_if_dvd_monic_vanishing_at {α: U → k} [DecidableEq k] {S: Finset U} {a: k[X]}:
  Set.InjOn α S →
  a ∣ monic_vanishing_at α S →
  a = C a.leadingCoeff * monic_vanishing_at α (S.filter (fun s ↦ a.eval (α s) = 0)) := by
  intro injective aA
  let T := S.filter (fun s ↦ a.eval (α s) = 0)
  let U := S.filter (fun s ↦ ¬a.eval (α s) = 0)
  let f := monic_vanishing_at α T
  let g := monic_vanishing_at α U
  have fg: f*g = monic_vanishing_at α S := monic_vanishing_at_filter α S (fun s ↦ a.eval (α s) = 0)
  have ag: IsCoprime a g := coprime_monic_vanishing_at_nonroots
  have afg: a ∣ f*g := by rw[fg]; exact aA
  have af: a ∣ f := dvd_if_dvd_mul_right ag afg
  have injt: Set.InjOn α T := by apply injon_filter α S _ injective
  have troots: ∀ s, s ∈ T → a.eval (α s) = 0 := by simp only [mem_filter, and_imp, imp_self, implies_true]
  have fa: f ∣ a := monic_vanishing_at_divides T a injt troots
  have fmonic: f.leadingCoeff = 1 := leadCoeff_monic_vanishing_at
  apply eq_lead_mul_if_dvd_dvd af fa fmonic

theorem lead_dvd {f g: k[X]}:
  f ∣ g → f.leadingCoeff ∣ g.leadingCoeff := by
  intro fg
  cases fg with | intro q gfq =>
    apply Dvd.intro q.leadingCoeff
    rw[←leadingCoeff_mul,gfq]

theorem lead_nonzero_if_dvd_monic_vanishing_at {α: U → k} [DecidableEq k] {S: Finset U} {a: k[X]}:
  a ∣ monic_vanishing_at α S →
  a.leadingCoeff ≠ 0 := by
  intro aA
  have Alead: (monic_vanishing_at α S).leadingCoeff = 1 := leadCoeff_monic_vanishing_at
  have aAlead: a.leadingCoeff ∣ (monic_vanishing_at α S).leadingCoeff := lead_dvd aA
  rw[Alead] at aAlead
  apply nonzero_if_dvd_one aAlead

theorem roots_card_eq_degree_if_dvd_monic_vanishing_at {α: U → k} [DecidableEq k] {S: Finset U} {a: k[X]}:
  Set.InjOn α S →
  a ∣ monic_vanishing_at α S →
  (S.filter (fun s ↦ a.eval (α s) = 0)).card = a.degree := by
  intro injective aA
  have acm: a = C a.leadingCoeff * monic_vanishing_at α (S.filter (fun s ↦ a.eval (α s) = 0)) := lead_mul_monic_vanishing_at_roots_if_dvd_monic_vanishing_at injective aA
  have leadnonzero: a.leadingCoeff ≠ 0 := lead_nonzero_if_dvd_monic_vanishing_at aA
  have deglead: (C a.leadingCoeff).degree = 0 := degree_C leadnonzero
  calc
    _ = (monic_vanishing_at α (S.filter (fun s ↦ a.eval (α s) = 0))).degree := (degree_monic_vanishing_at α (S.filter (fun s ↦ a.eval (α s) = 0))).symm
    _ = 0 + (monic_vanishing_at α (S.filter (fun s ↦ a.eval (α s) = 0))).degree := by group
    _ = (C a.leadingCoeff).degree + (monic_vanishing_at α (S.filter (fun s ↦ a.eval (α s) = 0))).degree := by rw[deglead]
    _ = (C a.leadingCoeff * monic_vanishing_at α (S.filter (fun s ↦ a.eval (α s) = 0))).degree := by rw[degree_mul]
    _ = a.degree := by rw[← acm]

theorem dvd_if_mul_C_nonzero_and_dvd {a b c: k[X]} {r: k}:
  a = C r * b → r ≠ 0 → b ∣ c → a ∣ c := by
  intro arb rnonzero bc
  cases bc with | intro q cbq =>
    use C (r⁻¹) * q
    calc
      c = b * q := cbq
      _ = 1 * (b * q) := by group
      _ = C 1 * (b * q) := by rw[map_one]
      _ = C (r⁻¹ * r) * (b * q) := by rw[inv_mul_cancel rnonzero]
      _ = C (r⁻¹) * C r * (b * q) := by rw[C_mul]
      _ = (C r * b) * (C (r⁻¹) * q) := by group
      _ = a * (C (r⁻¹) * q) := by rw[arb]

theorem dvd_if_dvd_monic_vanishing_at_and_root_imp_root {α: U → k} [DecidableEq k] {S: Finset U} {a b: k[X]}:
  Set.InjOn α S →
  a ∣ monic_vanishing_at α S →
  (∀ s, s ∈ S → a.eval (α s) = 0 → b.eval (α s) = 0) →
  a ∣ b :=
by
  intro injective aA ab
  have acm: a = C a.leadingCoeff * monic_vanishing_at α (S.filter (fun s ↦ a.eval (α s) = 0)) := lead_mul_monic_vanishing_at_roots_if_dvd_monic_vanishing_at injective aA
  have leadnonzero: a.leadingCoeff ≠ 0 := lead_nonzero_if_dvd_monic_vanishing_at aA
  have mb: monic_vanishing_at α (S.filter (fun s ↦ a.eval (α s) = 0)) ∣ b := by
    apply monic_vanishing_at_divides
    apply injon_filter α S _ injective
    simp
    apply ab
  exact dvd_if_mul_C_nonzero_and_dvd acm leadnonzero mb

theorem sum_C_mul_monic_vanishing_at_except_eq_interpolator {α: U → k} {S: Finset U} {A: k[X]} (c: U → k):
  Set.InjOn α S →
  A = monic_vanishing_at α S →
  ∑ s in S, C (c s) * monic_vanishing_at_except α S s
  = interpolator α S (fun s ↦ c s * A.diff.eval (α s)) :=
by
  intro injective Adef
  calc
    _ = ∑ s in S, C (c s * A.diff.eval (α s)) * vanishing_at_except_1_at α S s := by
          apply Finset.sum_congr
          simp only
          intro s sinS
          rw[Adef]
          unfold vanishing_at_except_1_at
          rw[derivative_monic_vanishing_at_eval sinS]
          rw[mul_comm (monic_vanishing_at_except α S s)]
          rw[←mul_assoc]
          rw[←C_mul]
          rw[mul_assoc]
          rw[mul_inv_cancel (monic_vanishing_at_except_eval_nonzero α S s injective sinS)]
          rw[mul_one]
    _ = interpolator α S (fun s ↦ c s * A.diff.eval (α s)) := rfl

end vanishing
