import Definitions.Def_ModularCurve_LaurentCoeff
import Definitions.Def_ModularCurve_JqCoeff

set_option autoImplicit false

noncomputable section

namespace ModularCurve
namespace CharPReduction

section IntegralCoeffs

variable {L : Type*} [CommRing L] (A : Subring L)

def integralCoeffs : Subring (LaurentSeries L) where
  carrier := {x | ∀ n : ℤ, x.coeff n ∈ A}
  zero_mem' := fun n => by simp
  one_mem' := fun n => by
    rw [HahnSeries.coeff_one]
    split_ifs
    · exact A.one_mem
    · exact A.zero_mem
  add_mem' := fun hx hy n => by
    rw [HahnSeries.coeff_add]
    exact A.add_mem (hx n) (hy n)
  neg_mem' := fun hx n => by
    rw [HahnSeries.coeff_neg]
    exact A.neg_mem (hx n)
  mul_mem' := fun hx hy n => by
    rw [HahnSeries.coeff_mul]
    exact A.sum_mem fun ij _ => A.mul_mem (hx ij.1) (hy ij.2)

theorem mem_integralCoeffs {x : LaurentSeries L} :
    x ∈ integralCoeffs A ↔ ∀ n : ℤ, x.coeff n ∈ A :=
  Iff.rfl

private def packCoeffs (x : integralCoeffs A) : LaurentSeries A where
  coeff n := ⟨(x : LaurentSeries L).coeff n, x.2 n⟩
  isPWO_support' := (x : LaurentSeries L).isPWO_support.mono
    (fun _ hn h => hn (Subtype.ext h))

private theorem coeffMap_subtype_packCoeffs (x : integralCoeffs A) :
    coeffMap A.subtype (packCoeffs A x) = (x : LaurentSeries L) := by
  ext n
  rfl

private theorem coeffMap_subtype_injective :
    Function.Injective (coeffMap A.subtype : LaurentSeries A →+* LaurentSeries L) := by
  intro x y h
  ext n
  exact congrArg (fun z : LaurentSeries L => z.coeff n) h

private def packCoeffsHom : integralCoeffs A →+* LaurentSeries A where
  toFun := packCoeffs A
  map_one' := coeffMap_subtype_injective A <| by
    rw [coeffMap_subtype_packCoeffs, map_one, OneMemClass.coe_one]
  map_mul' x y := coeffMap_subtype_injective A <| by
    rw [map_mul, coeffMap_subtype_packCoeffs, coeffMap_subtype_packCoeffs,
      coeffMap_subtype_packCoeffs, MulMemClass.coe_mul]
  map_zero' := coeffMap_subtype_injective A <| by
    rw [coeffMap_subtype_packCoeffs, map_zero, ZeroMemClass.coe_zero]
  map_add' x y := coeffMap_subtype_injective A <| by
    rw [map_add, coeffMap_subtype_packCoeffs, coeffMap_subtype_packCoeffs,
      coeffMap_subtype_packCoeffs, AddMemClass.coe_add]

variable {k : Type*} [CommRing k] (red : A →+* k)

def coeffRed : integralCoeffs A →+* LaurentSeries k :=
  (coeffMap red).comp (packCoeffsHom A)

@[simp]
theorem coeffRed_coeff (x : integralCoeffs A) (n : ℤ) :
    (coeffRed A red x).coeff n = red ⟨(x : LaurentSeries L).coeff n, x.2 n⟩ :=
  rfl

end IntegralCoeffs

section Reduction

variable {L : Type*} [CommRing L] {k : Type*} [Field k]
variable (A : Subring L) (red : A →+* k)
variable (R : Subring (LaurentSeries L)) (hR : R ≤ integralCoeffs A)

def redRes : R →+* LaurentSeries k :=
  (coeffRed A red).comp (Subring.inclusion hR)

theorem redRes_apply (r : R) :
    redRes A red R hR r = coeffRed A red ⟨(r : LaurentSeries L), hR r.2⟩ :=
  rfl

def redKer : Ideal R :=
  RingHom.ker (redRes A red R hR)

theorem redKer_isPrime : (redKer A red R hR).IsPrime :=
  RingHom.ker_isPrime _

theorem notMem_redKer_iff {s : R} :
    s ∉ redKer A red R hR ↔ redRes A red R hR s ≠ 0 := by
  rw [redKer, not_iff_not, RingHom.mem_ker]

theorem one_notMem_redKer : (1 : R) ∉ redKer A red R hR := by
  rw [notMem_redKer_iff, map_one]
  exact one_ne_zero

def localizedAtKer : Subring (LaurentSeries L) where
  carrier := {f | ∃ r s : R, s ∉ redKer A red R hR ∧ f * s = r}
  zero_mem' := ⟨0, 1, one_notMem_redKer A red R hR, by simp⟩
  one_mem' := ⟨1, 1, one_notMem_redKer A red R hR, by simp⟩
  add_mem' := by
    rintro f g ⟨r, s, hs, hf⟩ ⟨u, t, ht, hg⟩
    refine ⟨r * t + u * s, s * t, fun hmem =>
      (((redKer_isPrime A red R hR).mem_or_mem hmem).elim hs ht), ?_⟩
    push_cast
    calc (f + g) * ((s : LaurentSeries L) * t)
        = f * s * t + g * t * s := by ring
      _ = (r : LaurentSeries L) * t + u * s := by rw [hf, hg]
  neg_mem' := by
    rintro f ⟨r, s, hs, hf⟩
    exact ⟨-r, s, hs, by push_cast; rw [neg_mul, hf]⟩
  mul_mem' := by
    rintro f g ⟨r, s, hs, hf⟩ ⟨u, t, ht, hg⟩
    refine ⟨r * u, s * t, fun hmem =>
      (((redKer_isPrime A red R hR).mem_or_mem hmem).elim hs ht), ?_⟩
    push_cast
    calc f * g * ((s : LaurentSeries L) * t)
        = (f * s) * (g * t) := by ring
      _ = (r : LaurentSeries L) * u := by rw [hf, hg]

theorem mem_localizedAtKer {f : LaurentSeries L} :
    f ∈ localizedAtKer A red R hR ↔
      ∃ r s : R, s ∉ redKer A red R hR ∧ f * s = r :=
  Iff.rfl

theorem subring_le_localizedAtKer : R ≤ localizedAtKer A red R hR := fun r hr =>
  ⟨⟨r, hr⟩, 1, one_notMem_redKer A red R hR, by simp⟩

end Reduction

section LocalizedReduction

variable {L : Type*} [CommRing L] {k : Type*} [Field k]
variable {A : Subring L} {red : A →+* k}
variable {R : Subring (LaurentSeries L)} {hR : R ≤ integralCoeffs A}

open Classical in

def redLoc (f : localizedAtKer A red R hR) : LaurentSeries k :=
  redRes A red R hR (Classical.choose f.2) *
    (redRes A red R hR (Classical.choose (Classical.choose_spec f.2)))⁻¹

theorem redRes_ne_zero_of_notMem {s : R} (hs : s ∉ redKer A red R hR) :
    redRes A red R hR s ≠ 0 :=
  (notMem_redKer_iff A red R hR).mp hs

theorem redLoc_spec (f : localizedAtKer A red R hR) {r s : R}
    (hfs : (f : LaurentSeries L) * s = r) :
    redLoc f * redRes A red R hR s = redRes A red R hR r := by
  obtain ⟨hs₀, hfs₀⟩ := Classical.choose_spec (Classical.choose_spec f.2)
  have hs₀ne : redRes A red R hR (Classical.choose (Classical.choose_spec f.2)) ≠ 0 :=
    redRes_ne_zero_of_notMem (A := A) (red := red) (R := R) (hR := hR) hs₀
  have hcross : Classical.choose f.2 * s
      = r * Classical.choose (Classical.choose_spec f.2) := by
    have h : ((Classical.choose f.2 : R) : LaurentSeries L) * (s : LaurentSeries L)
        = (r : LaurentSeries L)
          * ((Classical.choose (Classical.choose_spec f.2) : R) : LaurentSeries L) := by
      linear_combination
        ((Classical.choose (Classical.choose_spec f.2) : R) : LaurentSeries L) * hfs
          - (s : LaurentSeries L) * hfs₀
    exact_mod_cast h
  have hred := congrArg (redRes A red R hR) hcross
  rw [map_mul, map_mul] at hred
  apply mul_right_cancel₀ hs₀ne
  calc redLoc f * redRes A red R hR s
        * redRes A red R hR (Classical.choose (Classical.choose_spec f.2))
      = redRes A red R hR (Classical.choose f.2) * redRes A red R hR s
        * ((redRes A red R hR (Classical.choose (Classical.choose_spec f.2)))⁻¹
          * redRes A red R hR (Classical.choose (Classical.choose_spec f.2))) := by
        rw [redLoc]; ring
    _ = redRes A red R hR (Classical.choose f.2) * redRes A red R hR s := by
        rw [inv_mul_cancel₀ hs₀ne, mul_one]
    _ = redRes A red R hR r
        * redRes A red R hR (Classical.choose (Classical.choose_spec f.2)) := hred

theorem redLoc_coe (r : R) :
    redLoc ⟨(r : LaurentSeries L), subring_le_localizedAtKer A red R hR r.2⟩
      = redRes A red R hR r := by
  have h := redLoc_spec (A := A) (red := red) (R := R) (hR := hR)
    ⟨(r : LaurentSeries L), subring_le_localizedAtKer A red R hR r.2⟩
    (r := r) (s := 1) (by simp)
  rwa [map_one, mul_one] at h

def redLocHom : localizedAtKer A red R hR →+* LaurentSeries k where
  toFun := redLoc
  map_one' := by
    have h := redLoc_spec (A := A) (red := red) (R := R) (hR := hR)
      (1 : localizedAtKer A red R hR) (r := 1) (s := 1) (by simp)
    rwa [map_one, mul_one] at h
  map_mul' f g := by
    obtain ⟨r, s, hs, hf⟩ := f.2
    obtain ⟨u, t, ht, hg⟩ := g.2
    have hsne := redRes_ne_zero_of_notMem (A := A) (red := red) (R := R) (hR := hR) hs
    have htne := redRes_ne_zero_of_notMem (A := A) (red := red) (R := R) (hR := hR) ht
    have hfg : ((f * g : localizedAtKer A red R hR) : LaurentSeries L)
        * ((s * t : R) : LaurentSeries L) = ((r * u : R) : LaurentSeries L) := by
      push_cast
      calc (f : LaurentSeries L) * g * ((s : LaurentSeries L) * t)
          = ((f : LaurentSeries L) * s) * ((g : LaurentSeries L) * t) := by ring
        _ = (r : LaurentSeries L) * u := by rw [hf, hg]
    have h2 := redLoc_spec (f * g) hfg
    have hfspec := redLoc_spec f hf
    have hgspec := redLoc_spec g hg
    apply mul_right_cancel₀ (mul_ne_zero hsne htne)
    calc redLoc (f * g) * (redRes A red R hR s * redRes A red R hR t)
        = redLoc (f * g) * redRes A red R hR (s * t) := by rw [map_mul]
      _ = redRes A red R hR (r * u) := h2
      _ = (redLoc f * redRes A red R hR s) * (redLoc g * redRes A red R hR t) := by
          rw [map_mul, hfspec, hgspec]
      _ = redLoc f * redLoc g * (redRes A red R hR s * redRes A red R hR t) := by ring
  map_zero' := by
    have h := redLoc_spec (A := A) (red := red) (R := R) (hR := hR)
      (0 : localizedAtKer A red R hR) (r := 0) (s := 1) (by simp)
    rwa [map_one, mul_one, map_zero] at h
  map_add' f g := by
    obtain ⟨r, s, hs, hf⟩ := f.2
    obtain ⟨u, t, ht, hg⟩ := g.2
    have hsne := redRes_ne_zero_of_notMem (A := A) (red := red) (R := R) (hR := hR) hs
    have htne := redRes_ne_zero_of_notMem (A := A) (red := red) (R := R) (hR := hR) ht
    have hfg : ((f + g : localizedAtKer A red R hR) : LaurentSeries L)
        * ((s * t : R) : LaurentSeries L) = ((r * t + u * s : R) : LaurentSeries L) := by
      push_cast
      calc ((f : LaurentSeries L) + g) * ((s : LaurentSeries L) * t)
          = ((f : LaurentSeries L) * s) * t + ((g : LaurentSeries L) * t) * s := by ring
        _ = (r : LaurentSeries L) * t + (u : LaurentSeries L) * s := by rw [hf, hg]
    have h2 := redLoc_spec (f + g) hfg
    have hfspec := redLoc_spec f hf
    have hgspec := redLoc_spec g hg
    apply mul_right_cancel₀ (mul_ne_zero hsne htne)
    calc redLoc (f + g) * (redRes A red R hR s * redRes A red R hR t)
        = redLoc (f + g) * redRes A red R hR (s * t) := by rw [map_mul]
      _ = redRes A red R hR (r * t + u * s) := h2
      _ = (redLoc f * redRes A red R hR s) * redRes A red R hR t
          + (redLoc g * redRes A red R hR t) * redRes A red R hR s := by
          rw [map_add, map_mul, map_mul, hfspec, hgspec]
      _ = (redLoc f + redLoc g) * (redRes A red R hR s * redRes A red R hR t) := by ring

@[simp]
theorem redLocHom_apply (f : localizedAtKer A red R hR) : redLocHom f = redLoc f :=
  rfl

end LocalizedReduction

section Modular

variable (N : ℕ) [NeZero N]
variable {L : Type*} [CommRing L] (A : Subring L)

def constSeries : A →+* LaurentSeries L :=
  (algebraMap L (LaurentSeries L)).comp A.subtype

def modularRing : Subring (LaurentSeries L) :=
  Subring.closure (Set.range (constSeries A) ∪ {jqModC L, jqNModC L N})

theorem jqModC_mem_modularRing : jqModC L ∈ modularRing N A :=
  Subring.subset_closure (Or.inr (Set.mem_insert _ _))

theorem jqNModC_mem_modularRing : jqNModC L N ∈ modularRing N A :=
  Subring.subset_closure (Or.inr (Set.mem_insert_of_mem _ rfl))

theorem constSeries_mem_modularRing (a : A) : constSeries A a ∈ modularRing N A :=
  Subring.subset_closure (Or.inl ⟨a, rfl⟩)

theorem jqModC_mem_integralCoeffs : jqModC L ∈ integralCoeffs A := by
  intro n
  rw [jqModC_eq_map_intCast, HahnSeries.map_coeff]
  exact intCast_mem A _

theorem jqNModC_mem_integralCoeffs : jqNModC L N ∈ integralCoeffs A := by
  intro n
  by_cases h : ((N : ℕ) : ℤ) ∣ n
  · obtain ⟨m, rfl⟩ := h
    rw [jqNModC, qExpand_coeff_mul]
    exact jqModC_mem_integralCoeffs A m
  · rw [jqNModC, qExpand_coeff_of_not_dvd _ _ h]
    exact A.zero_mem

theorem constSeries_mem_integralCoeffs (a : A) :
    constSeries A a ∈ integralCoeffs A := by
  intro n
  have : constSeries A a = HahnSeries.single 0 (a : L) := by
    show algebraMap L (LaurentSeries L) (a : L) = HahnSeries.single 0 (a : L)
    ext m
    rw [HahnSeries.algebraMap_apply', show algebraMap L (PowerSeries L) (a : L)
      = PowerSeries.C (a : L) from by simp, HahnSeries.ofPowerSeries_C]
    rfl
  rw [this, HahnSeries.coeff_single]
  split_ifs
  · exact a.2
  · exact A.zero_mem

theorem modularRing_le_integralCoeffs : modularRing N A ≤ integralCoeffs A := by
  rw [modularRing, Subring.closure_le]
  rintro x (⟨a, rfl⟩ | rfl | rfl)
  · exact constSeries_mem_integralCoeffs A a
  · exact jqModC_mem_integralCoeffs A
  · exact jqNModC_mem_integralCoeffs N A

variable {k : Type*} [Field k] (red : A →+* k)

def modularLocalized : Subring (LaurentSeries L) :=
  localizedAtKer A red (modularRing N A) (modularRing_le_integralCoeffs N A)

def modularRedLocHom : modularLocalized N A red →+* LaurentSeries k :=
  redLocHom (A := A) (red := red) (R := modularRing N A)
    (hR := modularRing_le_integralCoeffs N A)

end Modular

end CharPReduction
end ModularCurve

end