import Mathlib.AlgebraicGeometry.EllipticCurve.DivisionPolynomial.Basic
import Mathlib.Algebra.MvPolynomial.CommRing

open Polynomial
open scoped Polynomial.Bivariate

noncomputable section

namespace WeierstrassCurve

variable {R : Type*} [CommRing R] (W : WeierstrassCurve R)

def ψDbl (n : ℤ) : R[X][Y] :=
  complEDS₂ W.ψ₂ (C W.Ψ₃) (C W.preΨ₄) n

def twoω (n : ℤ) : R[X][Y] :=
  W.ψDbl n - W.ψ n * (C (C W.a₁) * W.φ n + C (C W.a₃) * W.ψ n ^ 2)

namespace Universal

abbrev Poly : Type := MvPolynomial (Fin 5) ℤ

def curve : WeierstrassCurve Poly :=
  ⟨MvPolynomial.X 0, MvPolynomial.X 1, MvPolynomial.X 2, MvPolynomial.X 3, MvPolynomial.X 4⟩

def specialize (W : WeierstrassCurve R) : Poly →+* R :=
  MvPolynomial.eval₂Hom (Int.castRingHom R) ![W.a₁, W.a₂, W.a₃, W.a₄, W.a₆]

@[simp] lemma specialize_X_zero : specialize W (MvPolynomial.X 0) = W.a₁ := by simp [specialize]
@[simp] lemma specialize_X_one : specialize W (MvPolynomial.X 1) = W.a₂ := by simp [specialize]
@[simp] lemma specialize_X_two : specialize W (MvPolynomial.X 2) = W.a₃ := by simp [specialize]
@[simp] lemma specialize_X_three : specialize W (MvPolynomial.X 3) = W.a₄ := by simp [specialize]
@[simp] lemma specialize_X_four : specialize W (MvPolynomial.X 4) = W.a₆ := by simp [specialize]

lemma map_specialize : curve.map (specialize W) = W := by
  rcases W with ⟨a₁, a₂, a₃, a₄, a₆⟩
  simp [curve, specialize, WeierstrassCurve.map]

def halveCoeff (r : Poly) : Poly :=
  AddMonoidAlgebra.ofCoeff
    (Finsupp.mapRange (fun z : ℤ => z / 2) (by simp) (AddMonoidAlgebra.coeff r))

@[simp] lemma halveCoeff_zero : halveCoeff 0 = 0 := by
  simp [halveCoeff]

lemma coeff_halveCoeff (r : Poly) (m : Fin 5 →₀ ℕ) :
    MvPolynomial.coeff m (halveCoeff r) = MvPolynomial.coeff m r / 2 :=
  Finsupp.mapRange_apply (hf := by simp) ..

def halveX (p : Poly[X]) : Poly[X] :=
  ⟨AddMonoidAlgebra.ofCoeff (p.toFinsupp.coeff.mapRange halveCoeff halveCoeff_zero)⟩

@[simp] lemma coeff_halveX (p : Poly[X]) (i : ℕ) : (halveX p).coeff i = halveCoeff (p.coeff i) := by
  rcases p with ⟨f⟩
  simp [halveX, Polynomial.coeff]

@[simp] lemma halveX_zero : halveX 0 = 0 := by
  ext i
  simp

def halve (p : Poly[X][Y]) : Poly[X][Y] :=
  ⟨AddMonoidAlgebra.ofCoeff (p.toFinsupp.coeff.mapRange halveX halveX_zero)⟩

@[simp] lemma coeff_halve (p : Poly[X][Y]) (i : ℕ) : (halve p).coeff i = halveX (p.coeff i) := by
  rcases p with ⟨f⟩
  simp [halve, Polynomial.coeff]

lemma halveCoeff_two_mul (r : Poly) : halveCoeff (2 * r) = r := by
  ext m
  rw [coeff_halveCoeff, show (2 : Poly) = MvPolynomial.C 2 from (map_ofNat MvPolynomial.C 2).symm,
    MvPolynomial.coeff_C_mul]
  exact Int.mul_ediv_cancel_left _ two_ne_zero

lemma halveX_two_mul (p : Poly[X]) : halveX (2 * p) = p := by
  ext i
  rw [coeff_halveX, show (2 : Poly[X]) = C 2 from (map_ofNat C 2).symm, coeff_C_mul,
    halveCoeff_two_mul]

lemma halve_two_mul (p : Poly[X][Y]) : halve (2 * p) = p := by
  ext i : 1
  rw [coeff_halve, show (2 : Poly[X][Y]) = C 2 from (map_ofNat C 2).symm, coeff_C_mul,
    halveX_two_mul]

end Universal

def ω (n : ℤ) : R[X][Y] :=
  (Universal.halve (Universal.curve.twoω n)).map (mapRingHom (Universal.specialize W))

lemma ψ_mul_ψDbl (n : ℤ) : W.ψ n * W.ψDbl n = W.ψ (2 * n) :=
  normEDS_mul_complEDS₂ ..

lemma ψDbl_zero : W.ψDbl 0 = 2 :=
  complEDS₂_zero ..

lemma ψDbl_one : W.ψDbl 1 = W.ψ₂ :=
  complEDS₂_one ..

lemma ψDbl_neg (n : ℤ) : W.ψDbl (-n) = W.ψDbl n :=
  complEDS₂_neg ..

lemma twoω_zero : W.twoω 0 = 2 := by
  rw [twoω, ψDbl_zero, ψ_zero]
  ring

lemma twoω_one : W.twoω 1 = 2 * Y := by
  rw [twoω, ψDbl_one, ψ_one, φ_one, ψ₂, Affine.polynomialY]
  simp only [map_add, map_mul, map_ofNat]
  ring

lemma twoω_neg (n : ℤ) :
    W.twoω (-n) = W.twoω n + 2 * (W.ψ n * (C (C W.a₁) * W.φ n + C (C W.a₃) * W.ψ n ^ 2)) := by
  rw [twoω, twoω, ψDbl_neg, ψ_neg, φ_neg]
  ring

section Map

variable {S : Type*} [CommRing S] (f : R →+* S)

lemma map_ψDbl (n : ℤ) : (W.map f).ψDbl n = (W.ψDbl n).map (mapRingHom f) := by
  rw [ψDbl, ψDbl, ← coe_mapRingHom, map_complEDS₂, map_ψ₂, map_Ψ₃, map_preΨ₄, coe_mapRingHom,
    map_C, map_C, coe_mapRingHom]

lemma map_twoω (n : ℤ) : (W.map f).twoω n = (W.twoω n).map (mapRingHom f) := by
  simp only [twoω, map_ψDbl, map_ψ, map_φ, map_a₁, map_a₃, Polynomial.map_sub, Polynomial.map_mul,
    Polynomial.map_add, Polynomial.map_pow, map_C, coe_mapRingHom]

end Map

lemma ω_eq_map_of_twoω_eq {n : ℤ} {q : Universal.Poly[X][Y]}
    (hq : Universal.curve.twoω n = 2 * q) :
    W.ω n = q.map (mapRingHom (Universal.specialize W)) := by
  rw [ω, hq, Universal.halve_two_mul]

lemma ω_one : W.ω 1 = Y := by
  rw [ω, twoω_one, Universal.halve_two_mul, Polynomial.map_X]

theorem two_mul_ω_of_twoω_universal_eq {n : ℤ} {q : Universal.Poly[X][Y]}
    (hq : Universal.curve.twoω n = 2 * q) : 2 * W.ω n = W.twoω n := by
  calc 2 * W.ω n = 2 * q.map (mapRingHom (Universal.specialize W)) := by
        rw [W.ω_eq_map_of_twoω_eq hq]
    _ = (2 * q).map (mapRingHom (Universal.specialize W)) := by
        rw [Polynomial.map_mul, Polynomial.map_ofNat]
    _ = (Universal.curve.map (Universal.specialize W)).twoω n := by rw [← hq, map_twoω]
    _ = W.twoω n := by rw [Universal.map_specialize]

end WeierstrassCurve

end