import Mathlib

open Matrix Matrix.SpecialLinearGroup OnePoint CongruenceSubgroup

open scoped MatrixGroups

namespace ModularCurve

noncomputable section

def ratPoint (a c : ℤ) : OnePoint ℚ :=
  if c = 0 then ∞ else (((a : ℚ) / (c : ℚ) : ℚ) : OnePoint ℚ)

@[simp] lemma ratPoint_zero_right (a : ℤ) : ratPoint a 0 = ∞ := if_pos rfl

lemma ratPoint_of_ne_zero (a : ℤ) {c : ℤ} (hc : c ≠ 0) :
    ratPoint a c = (((a : ℚ) / (c : ℚ) : ℚ) : OnePoint ℚ) := if_neg hc

@[simp] lemma ratPoint_one_zero : ratPoint 1 0 = ∞ := if_pos rfl

@[simp] lemma ratPoint_zero_one : ratPoint 0 1 = ((0 : ℚ) : OnePoint ℚ) := by
  simp [ratPoint]

lemma ratPoint_eq_infty_iff {a c : ℤ} : ratPoint a c = ∞ ↔ c = 0 := by
  rcases eq_or_ne c 0 with rfl | hc
  · simp
  · simp [ratPoint_of_ne_zero a hc, hc]

lemma ratPoint_mul_left {l : ℤ} (hl : l ≠ 0) (a c : ℤ) :
    ratPoint (l * a) (l * c) = ratPoint a c := by
  rcases eq_or_ne c 0 with rfl | hc
  · simp
  · rw [ratPoint_of_ne_zero _ (mul_ne_zero hl hc), ratPoint_of_ne_zero _ hc]
    push_cast
    rw [mul_div_mul_left _ _ (by exact_mod_cast hl)]

lemma ratPoint_neg_neg (a c : ℤ) : ratPoint (-a) (-c) = ratPoint a c := by
  simpa using ratPoint_mul_left (l := -1) (by norm_num) a c

lemma exists_ratPoint_eq (x : OnePoint ℚ) : ∃ a c : ℤ, IsCoprime a c ∧ ratPoint a c = x := by
  cases x with
  | infty => exact ⟨1, 0, isCoprime_one_left, rfl⟩
  | coe q =>
    refine ⟨q.num, (q.den : ℤ), ?_, ?_⟩
    · rw [Int.isCoprime_iff_gcd_eq_one]
      simpa [Int.gcd] using q.reduced
    · rw [ratPoint_of_ne_zero _ (by exact_mod_cast q.den_nz), OnePoint.coe_eq_coe,
        show ((q.den : ℤ) : ℚ) = (q.den : ℚ) by push_cast; ring, Rat.num_div_den]

@[simp] lemma mapGL_apply_coe (γ : SL(2, ℤ)) (i j : Fin 2) :
    (mapGL ℚ γ) i j = ((γ i j : ℤ) : ℚ) := by
  simp [mapGL]

lemma mapGL_smul_infty (γ : SL(2, ℤ)) :
    mapGL ℚ γ • (∞ : OnePoint ℚ) = ratPoint (γ 0 0) (γ 1 0) := by
  rw [smul_infty_eq_ite, mapGL_apply_coe, mapGL_apply_coe]
  rcases eq_or_ne (γ 1 0) 0 with h | h
  · simp [h]
  · rw [if_neg (by exact_mod_cast h), ratPoint_of_ne_zero _ h]

lemma mapGL_smul_ratPoint (γ : SL(2, ℤ)) {a c : ℤ} (h : a ≠ 0 ∨ c ≠ 0) :
    mapGL ℚ γ • ratPoint a c =
      ratPoint (γ 0 0 * a + γ 0 1 * c) (γ 1 0 * a + γ 1 1 * c) := by
  rcases eq_or_ne c 0 with rfl | hc
  ·
    replace h : a ≠ 0 := h.resolve_right (by simp)
    rw [ratPoint_zero_right, mapGL_smul_infty]
    simp only [mul_zero, add_zero]
    rw [mul_comm (γ 0 0) a, mul_comm (γ 1 0) a, ratPoint_mul_left h]
  ·
    rw [ratPoint_of_ne_zero _ hc, smul_some_eq_ite]
    have hcQ : (c : ℚ) ≠ 0 := by exact_mod_cast hc
    have hden : (mapGL ℚ γ) 1 0 * ((a : ℚ) / c) + (mapGL ℚ γ) 1 1
        = (((γ 1 0 * a + γ 1 1 * c : ℤ) : ℚ)) / (c : ℚ) := by
      simp only [mapGL_apply_coe]
      field_simp
      push_cast
      ring
    have hnum : (mapGL ℚ γ) 0 0 * ((a : ℚ) / c) + (mapGL ℚ γ) 0 1
        = (((γ 0 0 * a + γ 0 1 * c : ℤ) : ℚ)) / (c : ℚ) := by
      simp only [mapGL_apply_coe]
      field_simp
      push_cast
      ring
    rcases eq_or_ne (γ 1 0 * a + γ 1 1 * c) 0 with h0 | h0
    · rw [if_pos (by rw [hden, h0]; simp), h0, ratPoint_zero_right]
    · rw [if_neg (by rw [hden]; exact div_ne_zero (by exact_mod_cast h0) hcQ),
        ratPoint_of_ne_zero _ h0, hnum, hden, OnePoint.coe_eq_coe,
        div_div_div_cancel_right₀ hcQ]

lemma isCoprime_smul_pair (γ : SL(2, ℤ)) {a c : ℤ} (h : IsCoprime a c) :
    IsCoprime (γ 0 0 * a + γ 0 1 * c) (γ 1 0 * a + γ 1 1 * c) := by
  have := IsCoprime.mulVecSL (v := ![a, c]) (by simpa using h) γ
  simpa [Matrix.mulVec_fin_two] using this

variable (N : ℕ)

def Gamma0Q : Subgroup (GL (Fin 2) ℚ) := (Gamma0 N).map (mapGL ℚ)

lemma mem_Gamma0Q {g : GL (Fin 2) ℚ} :
    g ∈ Gamma0Q N ↔ ∃ γ ∈ Gamma0 N, mapGL ℚ γ = g := Iff.rfl

abbrev CuspSpace : Type := MulAction.orbitRel.Quotient (Gamma0Q N) (OnePoint ℚ)

namespace CuspSpace

def mk (x : OnePoint ℚ) : CuspSpace N := ⟦x⟧

lemma mk_surjective : Function.Surjective (mk N) := by
  rintro ⟨x⟩; exact ⟨x, rfl⟩

variable {N}

lemma mk_eq_mk_iff {x y : OnePoint ℚ} :
    mk N x = mk N y ↔ ∃ γ ∈ Gamma0 N, mapGL ℚ γ • y = x := by
  rw [mk, mk, Quotient.eq, MulAction.orbitRel_apply, MulAction.mem_orbit_iff]
  constructor
  · rintro ⟨⟨g, γ, hγ, rfl⟩, hg⟩
    exact ⟨γ, hγ, hg⟩
  · rintro ⟨γ, hγ, hg⟩
    exact ⟨⟨mapGL ℚ γ, γ, hγ, rfl⟩, hg⟩

lemma mk_smul {γ : SL(2, ℤ)} (hγ : γ ∈ Gamma0 N) (x : OnePoint ℚ) :
    mk N (mapGL ℚ γ • x) = mk N x :=
  mk_eq_mk_iff.mpr ⟨γ, hγ, rfl⟩

variable (N)

lemma mk_smul_infty_eq_of_mem {a b : SL(2, ℤ)} (hab : a⁻¹ * b ∈ Gamma0 N) :
    mk N (mapGL ℚ a⁻¹ • ∞) = mk N (mapGL ℚ b⁻¹ • ∞) := by
  have key : mapGL ℚ b⁻¹ • (∞ : OnePoint ℚ)
      = mapGL ℚ (a⁻¹ * b)⁻¹ • (mapGL ℚ a⁻¹ • (∞ : OnePoint ℚ)) := by
    have h1 : (a⁻¹ * b)⁻¹ * a⁻¹ = b⁻¹ := by group
    rw [← mul_smul, ← map_mul, h1]
  rw [key]
  exact (mk_smul (inv_mem hab) _).symm

def fromCoset : SL(2, ℤ) ⧸ Gamma0 N → CuspSpace N :=
  Quotient.lift (fun g => mk N (mapGL ℚ g⁻¹ • ∞))
    (fun _ _ hab => mk_smul_infty_eq_of_mem N (QuotientGroup.leftRel_apply.mp hab))

@[simp] lemma fromCoset_mk (g : SL(2, ℤ)) :
    fromCoset N ((g : SL(2, ℤ) ⧸ Gamma0 N)) = mk N (mapGL ℚ g⁻¹ • ∞) := rfl

lemma fromCoset_surjective : Function.Surjective (fromCoset N) := by
  rintro ⟨x⟩
  obtain ⟨g, hg⟩ := OnePoint.exists_mem_SL2 ℤ x
  exact ⟨(g⁻¹ : SL(2, ℤ)), by rw [fromCoset_mk, inv_inv, hg]; rfl⟩

instance instFinite [NeZero N] : Finite (CuspSpace N) :=
  Finite.of_surjective _ (fromCoset_surjective N)

instance instNonempty : Nonempty (CuspSpace N) := ⟨mk N ∞⟩

instance instSubsingletonOfOne : Subsingleton (CuspSpace 1) := by
  constructor
  rintro ⟨x⟩ ⟨y⟩
  obtain ⟨g, hg⟩ := OnePoint.exists_mem_SL2 ℤ x
  obtain ⟨h, hh⟩ := OnePoint.exists_mem_SL2 ℤ y
  show mk 1 x = mk 1 y
  have hmem : ∀ γ : SL(2, ℤ), γ ∈ Gamma0 1 := fun γ => by
    rw [Gamma0_mem]; exact Subsingleton.elim _ _
  rw [← hg, ← hh]
  symm
  rw [show (mapGL ℚ h • (∞ : OnePoint ℚ)) = mapGL ℚ (h * g⁻¹) • (mapGL ℚ g • (∞ : OnePoint ℚ))
    by rw [← mul_smul, ← map_mul, inv_mul_cancel_right]]
  exact mk_smul (hmem _) _

def cuspDenomAux (x : OnePoint ℚ) : ℕ :=
  x.elim N fun q => Nat.gcd q.den N

@[simp] lemma cuspDenomAux_infty : cuspDenomAux N ∞ = N := rfl

@[simp] lemma cuspDenomAux_coe (q : ℚ) : cuspDenomAux N (q : OnePoint ℚ) = Nat.gcd q.den N :=
  rfl

variable {N}

lemma cuspDenomAux_ratPoint {a c : ℤ} (h : IsCoprime a c) :
    cuspDenomAux N (ratPoint a c) = Int.gcd c N := by
  rcases eq_or_ne c 0 with rfl | hc
  · simp [Int.gcd]
  · rw [ratPoint_of_ne_zero _ hc, cuspDenomAux_coe]
    have hcop : Nat.Coprime a.natAbs c.natAbs := Int.isCoprime_iff_gcd_eq_one.mp h
    have hden : (((a : ℚ) / (c : ℚ)).den : ℤ) = (c.natAbs : ℤ) := by
      rcases lt_or_gt_of_ne hc with hneg | hpos
      · have : ((a : ℚ) / (c : ℚ)) = (((-a : ℤ) : ℚ) / ((-c : ℤ) : ℚ)) := by
          push_cast; rw [neg_div_neg_eq]
        rw [this, Rat.den_div_eq_of_coprime (by omega) (by simpa using hcop)]
        omega
      · rw [Rat.den_div_eq_of_coprime hpos hcop]
        omega
    have : ((a : ℚ) / (c : ℚ)).den = c.natAbs := by exact_mod_cast hden
    rw [this]
    rfl

lemma _root_.CongruenceSubgroup.Gamma0_mem_iff_dvd {N : ℕ} {γ : SL(2, ℤ)} :
    γ ∈ Gamma0 N ↔ (N : ℤ) ∣ γ 1 0 := by
  rw [Gamma0_mem, ZMod.intCast_zmod_eq_zero_iff_dvd]

lemma cuspDenomAux_dvd_smul {γ : SL(2, ℤ)} (hγ : γ ∈ Gamma0 N) (x : OnePoint ℚ) :
    cuspDenomAux N x ∣ cuspDenomAux N (mapGL ℚ γ • x) := by
  obtain ⟨a, c, hac, rfl⟩ := exists_ratPoint_eq x
  have hac' : a ≠ 0 ∨ c ≠ 0 := hac.ne_zero_or_ne_zero
  rw [mapGL_smul_ratPoint γ hac', cuspDenomAux_ratPoint hac,
    cuspDenomAux_ratPoint (isCoprime_smul_pair γ hac)]

  have h1 : (Int.gcd c N : ℤ) ∣ γ 1 0 * a + γ 1 1 * c := by
    refine dvd_add (Dvd.dvd.mul_right ?_ a) (Dvd.dvd.mul_left (Int.gcd_dvd_left ..) _)
    exact (Int.gcd_dvd_right ..).trans (Gamma0_mem_iff_dvd.mp hγ)
  exact Int.dvd_gcd h1 (Int.gcd_dvd_right ..)

lemma cuspDenomAux_smul {γ : SL(2, ℤ)} (hγ : γ ∈ Gamma0 N) (x : OnePoint ℚ) :
    cuspDenomAux N (mapGL ℚ γ • x) = cuspDenomAux N x := by
  refine Nat.dvd_antisymm ?_ (cuspDenomAux_dvd_smul hγ x)
  have := cuspDenomAux_dvd_smul (inv_mem hγ) (mapGL ℚ γ • x)
  rwa [← mul_smul, ← map_mul, inv_mul_cancel, map_one, one_smul] at this

variable (N)

def cuspDenom : CuspSpace N → ℕ :=
  Quotient.lift (cuspDenomAux N) (by
    rintro x y ⟨⟨g, γ, hγ, rfl⟩, rfl⟩
    exact cuspDenomAux_smul hγ y)

@[simp] lemma cuspDenom_mk (x : OnePoint ℚ) : cuspDenom N (mk N x) = cuspDenomAux N x := rfl

variable {N}

lemma cuspDenom_mk_ratPoint {a c : ℤ} (h : IsCoprime a c) :
    cuspDenom N (mk N (ratPoint a c)) = Int.gcd c N :=
  cuspDenomAux_ratPoint h

@[simp] lemma cuspDenom_mk_infty : cuspDenom N (mk N ∞) = N := rfl

@[simp] lemma cuspDenom_mk_zero : cuspDenom N (mk N ((0 : ℚ) : OnePoint ℚ)) = 1 := by
  rw [← ratPoint_zero_one, cuspDenom_mk_ratPoint (isCoprime_zero_left.mpr isUnit_one)]
  simp [Int.gcd]

lemma cuspDenom_dvd (x : CuspSpace N) : cuspDenom N x ∣ N := by
  obtain ⟨y, rfl⟩ := mk_surjective N x
  obtain ⟨a, c, hac, rfl⟩ := exists_ratPoint_eq y
  rw [cuspDenom_mk_ratPoint hac]
  simpa using Int.gcd_dvd_natAbs_right c (N : ℤ)

lemma cuspDenom_pos [NeZero N] (x : CuspSpace N) : 0 < cuspDenom N x :=
  Nat.pos_of_dvd_of_pos (cuspDenom_dvd x) (Nat.pos_of_ne_zero (NeZero.ne N))

lemma cuspDenom_mk_ratPoint_one_dvd {e : ℕ} (he : e ∣ N) :
    cuspDenom N (mk N (ratPoint 1 e)) = e := by
  rw [cuspDenom_mk_ratPoint isCoprime_one_left, Int.gcd_natCast_natCast, Nat.gcd_eq_left he]

def NormalFormCriterion (N : ℕ) : Prop :=
  ∀ (a a' : ℤ) (e : ℕ), e ∣ N → IsCoprime a (e : ℤ) → IsCoprime a' (e : ℤ) →
    (mk N (ratPoint a e) = mk N (ratPoint a' e) ↔
      (a : ZMod (Nat.gcd e (N / e))) = (a' : ZMod (Nat.gcd e (N / e))))

def Classification (N : ℕ) : Prop :=
  Nonempty (CuspSpace N ≃ (d : N.divisors) × (ZMod (Nat.gcd d (N / (d : ℕ))))ˣ)

end CuspSpace

end

end ModularCurve

open Matrix Matrix.SpecialLinearGroup OnePoint CongruenceSubgroup

open scoped MatrixGroups

namespace ModularCurve

noncomputable section

theorem gcd_sq_eq_mul_gcd_div {N d : ℕ} (hd : d ∣ N) :
    Nat.gcd (d ^ 2) N = d * Nat.gcd d (N / d) := by
  conv_lhs => rw [pow_two, ← Nat.mul_div_cancel' hd]
  exact Nat.gcd_mul_left d d (N / d)

theorem gcd_sq_gcd_left (c N : ℕ) : Nat.gcd (Nat.gcd c N ^ 2) N = Nat.gcd (c ^ 2) N := by
  rcases Nat.eq_zero_or_pos N with rfl | hN
  · simp [Nat.gcd_zero_right]
  rcases Nat.eq_zero_or_pos c with rfl | hc
  · rw [Nat.gcd_zero_left, zero_pow two_ne_zero, Nat.gcd_zero_left,
      Nat.gcd_eq_right (dvd_pow_self N two_ne_zero)]

  set g : ℕ := Nat.gcd c N with hg
  have hg0 : 0 < g := Nat.gcd_pos_of_pos_left N hc
  have hgc : g ∣ c := Nat.gcd_dvd_left c N
  have hgN : g ∣ N := Nat.gcd_dvd_right c N
  have hcop : Nat.Coprime (c / g) (N / g) := Nat.coprime_div_gcd_div_gcd hg0

  have hL : Nat.gcd (g ^ 2) N = g * Nat.gcd g (N / g) := gcd_sq_eq_mul_gcd_div hgN

  have hR : Nat.gcd (c ^ 2) N = g * Nat.gcd g (N / g) := by
    have h1 : c ^ 2 = g * (c / g * c) := by
      rw [← mul_assoc, Nat.mul_div_cancel' hgc, ← pow_two]
    have h2 : N = g * (N / g) := (Nat.mul_div_cancel' hgN).symm
    rw [h1]
    conv_lhs => rw [h2]
    rw [Nat.gcd_mul_left]
    congr 1
    rw [Nat.Coprime.gcd_mul_left_cancel c hcop]
    conv_lhs => rw [← Nat.div_mul_cancel hgc]
    rw [Nat.Coprime.gcd_mul_left_cancel g hcop]
  rw [hL, hR]

def cuspWidthOfDenom (N d : ℕ) : ℕ := N / Nat.gcd (d ^ 2) N

theorem gcd_mul_cuspWidthOfDenom (N d : ℕ) :
    Nat.gcd (d ^ 2) N * cuspWidthOfDenom N d = N :=
  Nat.mul_div_cancel' (Nat.gcd_dvd_right _ _)

theorem cuspWidthOfDenom_dvd (N d : ℕ) : cuspWidthOfDenom N d ∣ N :=
  Nat.div_dvd_of_dvd (Nat.gcd_dvd_right _ _)

theorem cuspWidthOfDenom_pos {N : ℕ} (hN : N ≠ 0) (d : ℕ) : 0 < cuspWidthOfDenom N d :=
  Nat.div_pos (Nat.le_of_dvd (Nat.pos_of_ne_zero hN) (Nat.gcd_dvd_right _ _))
    (Nat.gcd_pos_of_pos_right _ (Nat.pos_of_ne_zero hN))

theorem cuspWidthOfDenom_eq_div_div {N d : ℕ} (hd : d ∣ N) :
    cuspWidthOfDenom N d = N / d / Nat.gcd d (N / d) := by
  unfold cuspWidthOfDenom
  rw [gcd_sq_eq_mul_gcd_div hd, Nat.div_div_eq_div_mul]

theorem cuspWidthOfDenom_gcd (N d : ℕ) :
    cuspWidthOfDenom N (Nat.gcd d N) = cuspWidthOfDenom N d := by
  unfold cuspWidthOfDenom
  rw [gcd_sq_gcd_left]

theorem cuspWidthOfDenom_self {N : ℕ} (hN : N ≠ 0) : cuspWidthOfDenom N N = 1 := by
  unfold cuspWidthOfDenom
  rw [Nat.gcd_eq_right (dvd_pow_self N two_ne_zero), Nat.div_self (Nat.pos_of_ne_zero hN)]

theorem cuspWidthOfDenom_one (N : ℕ) : cuspWidthOfDenom N 1 = N := by
  unfold cuspWidthOfDenom
  rw [one_pow, Nat.gcd_one_left, Nat.div_one]

namespace CuspSpace

variable {N : ℕ}

def cuspWidth (x : CuspSpace N) : ℕ := cuspWidthOfDenom N (cuspDenom N x)

theorem cuspWidth_mk_ratPoint {a c : ℤ} (h : IsCoprime a c) :
    cuspWidth (mk N (ratPoint a c)) = N / Nat.gcd (c.natAbs ^ 2) N := by
  have h1 : cuspDenom N (mk N (ratPoint a c)) = Nat.gcd c.natAbs N := by
    rw [cuspDenom_mk_ratPoint h]
    show Nat.gcd c.natAbs (N : ℤ).natAbs = _
    rw [Int.natAbs_natCast]
  show cuspWidthOfDenom N (cuspDenom N (mk N (ratPoint a c))) = _
  rw [h1]
  show N / Nat.gcd (Nat.gcd c.natAbs N ^ 2) N = _
  rw [gcd_sq_gcd_left]

theorem cuspWidth_mk_infty (hN : N ≠ 0) : cuspWidth (mk N ∞) = 1 := by
  show cuspWidthOfDenom N (cuspDenom N (mk N ∞)) = 1
  rw [cuspDenom_mk_infty]
  exact cuspWidthOfDenom_self hN

theorem cuspWidth_mk_zero : cuspWidth (mk N ((0 : ℚ) : OnePoint ℚ)) = N := by
  show cuspWidthOfDenom N (cuspDenom N (mk N ((0 : ℚ) : OnePoint ℚ))) = N
  rw [cuspDenom_mk_zero]
  exact cuspWidthOfDenom_one N

theorem cuspWidth_dvd (x : CuspSpace N) : cuspWidth x ∣ N := cuspWidthOfDenom_dvd _ _

theorem cuspWidth_pos (hN : N ≠ 0) (x : CuspSpace N) : 0 < cuspWidth x :=
  cuspWidthOfDenom_pos hN _

end CuspSpace

end

end ModularCurve