import Mathlib
import Definitions.Def_CuspForm_HeckeAlgebra
import Definitions.Def_HeckeGalois_EichlerShimura

set_option autoImplicit false

noncomputable section

open UpperHalfPlane

open scoped MatrixGroups

namespace ModularCurve

section Period

def clamp01 (t : ℝ) : ℝ := max 0 (min t 1)

theorem clamp01_nonneg (t : ℝ) : 0 ≤ clamp01 t := le_max_left _ _

theorem clamp01_le_one (t : ℝ) : clamp01 t ≤ 1 := max_le zero_le_one (min_le_right _ _)

theorem continuous_clamp01 : Continuous clamp01 :=
  continuous_const.max (continuous_id.min continuous_const)

theorem clamp01_of_mem {t : ℝ} (ht : t ∈ Set.Icc (0 : ℝ) 1) : clamp01 t = t := by
  rw [clamp01, min_eq_left ht.2, max_eq_right ht.1]

def segmentPoint (τ₀ τ₁ : ℍ) (t : ℝ) : ℂ :=
  (1 - clamp01 t) • (τ₀ : ℂ) + clamp01 t • (τ₁ : ℂ)

theorem segmentPoint_im_pos (τ₀ τ₁ : ℍ) (t : ℝ) : 0 < (segmentPoint τ₀ τ₁ t).im := by
  have h0 := clamp01_nonneg t
  have h1 := clamp01_le_one t
  simp only [segmentPoint, Complex.add_im, Complex.smul_im, smul_eq_mul]
  rcases eq_or_lt_of_le h1 with h | h
  · rw [h]; simpa using τ₁.im_pos
  · exact add_pos_of_pos_of_nonneg (mul_pos (by linarith) τ₀.im_pos)
      (mul_nonneg h0 τ₁.im_pos.le)

theorem continuous_segmentPoint (τ₀ τ₁ : ℍ) : Continuous (segmentPoint τ₀ τ₁) :=
  ((continuous_const.sub continuous_clamp01).smul continuous_const).add
    (continuous_clamp01.smul continuous_const)

def segmentPath (τ₀ τ₁ : ℍ) (t : ℝ) : ℍ :=
  UpperHalfPlane.mk (segmentPoint τ₀ τ₁ t) (segmentPoint_im_pos τ₀ τ₁ t)

@[simp] theorem coe_segmentPath (τ₀ τ₁ : ℍ) (t : ℝ) :
    ((segmentPath τ₀ τ₁ t : ℍ) : ℂ) = segmentPoint τ₀ τ₁ t := rfl

theorem continuous_segmentPath (τ₀ τ₁ : ℍ) : Continuous (segmentPath τ₀ τ₁) :=
  (continuous_segmentPoint τ₀ τ₁).upperHalfPlaneMk _

variable (N : ℕ)

def periodIntegrand (τ₀ τ₁ : ℍ) (f : CuspForm (CongruenceSubgroup.Gamma0 N) 2) (t : ℝ) : ℂ :=
  f (segmentPath τ₀ τ₁ t) * ((τ₁ : ℂ) - τ₀)

theorem continuous_periodIntegrand (τ₀ τ₁ : ℍ) (f : CuspForm (CongruenceSubgroup.Gamma0 N) 2) :
    Continuous (periodIntegrand N τ₀ τ₁ f) :=
  ((f.holo'.continuous).comp (continuous_segmentPath τ₀ τ₁)).mul continuous_const

theorem intervalIntegrable_periodIntegrand (τ₀ τ₁ : ℍ)
    (f : CuspForm (CongruenceSubgroup.Gamma0 N) 2) (a b : ℝ) :
    IntervalIntegrable (periodIntegrand N τ₀ τ₁ f) MeasureTheory.volume a b :=
  (continuous_periodIntegrand N τ₀ τ₁ f).intervalIntegrable a b

theorem periodIntegrand_add (τ₀ τ₁ : ℍ) (f g : CuspForm (CongruenceSubgroup.Gamma0 N) 2) :
    periodIntegrand N τ₀ τ₁ (f + g) = periodIntegrand N τ₀ τ₁ f + periodIntegrand N τ₀ τ₁ g := by
  funext t
  simp [periodIntegrand, add_mul]

theorem periodIntegrand_smul (τ₀ τ₁ : ℍ) (c : ℂ) (f : CuspForm (CongruenceSubgroup.Gamma0 N) 2) :
    periodIntegrand N τ₀ τ₁ (c • f) = fun t => c * periodIntegrand N τ₀ τ₁ f t := by
  funext t
  simp [periodIntegrand, mul_assoc]

def periodAlong (τ₀ τ₁ : ℍ) : Module.Dual ℂ (CuspForm (CongruenceSubgroup.Gamma0 N) 2) where
  toFun f := ∫ t in (0 : ℝ)..1, periodIntegrand N τ₀ τ₁ f t
  map_add' f g := by
    rw [periodIntegrand_add]
    exact intervalIntegral.integral_add (intervalIntegrable_periodIntegrand N τ₀ τ₁ f 0 1)
      (intervalIntegrable_periodIntegrand N τ₀ τ₁ g 0 1)
  map_smul' c f := by
    rw [periodIntegrand_smul, RingHom.id_apply, smul_eq_mul]
    exact intervalIntegral.integral_const_mul c _

theorem periodAlong_apply (τ₀ τ₁ : ℍ) (f : CuspForm (CongruenceSubgroup.Gamma0 N) 2) :
    periodAlong N τ₀ τ₁ f = ∫ t in (0 : ℝ)..1, f (segmentPath τ₀ τ₁ t) * ((τ₁ : ℂ) - τ₀) :=
  rfl

def period (γ : CongruenceSubgroup.Gamma0 N) :
    Module.Dual ℂ (CuspForm (CongruenceSubgroup.Gamma0 N) 2) :=
  periodAlong N UpperHalfPlane.I ((γ : SL(2, ℤ)) • UpperHalfPlane.I)

theorem period_apply (γ : CongruenceSubgroup.Gamma0 N) (f : CuspForm (CongruenceSubgroup.Gamma0 N) 2) :
    period N γ f =
      ∫ t in (0 : ℝ)..1, f (segmentPath UpperHalfPlane.I ((γ : SL(2, ℤ)) • UpperHalfPlane.I) t) *
        ((((γ : SL(2, ℤ)) • UpperHalfPlane.I : ℍ) : ℂ) - (UpperHalfPlane.I : ℂ)) :=
  rfl

def periodLattice : Submodule ℤ (Module.Dual ℂ (CuspForm (CongruenceSubgroup.Gamma0 N) 2)) :=
  Submodule.span ℤ (Set.range (period N))

theorem period_mem_periodLattice (γ : CongruenceSubgroup.Gamma0 N) : period N γ ∈ periodLattice N :=
  Submodule.subset_span (Set.mem_range_self γ)

end Period

section Hecke

variable (N : ℕ) [NeZero N]

def cuspHeckeGen (ℓ : Nat.Primes) : CuspForm.heckeAlgebra N 2 (∅ : Set ℕ) :=
  if h : (ℓ : ℕ) ∣ N then CuspForm.heckeAlgebra.U (S := (∅ : Set ℕ)) ℓ.prop h (Set.notMem_empty _)
  else CuspForm.heckeAlgebra.T (S := (∅ : Set ℕ)) ℓ.prop h (Set.notMem_empty _)

theorem cuspHeckeGen_of_dvd (ℓ : Nat.Primes) (h : (ℓ : ℕ) ∣ N) :
    (cuspHeckeGen N ℓ : Module.End ℂ (CuspForm (CongruenceSubgroup.Gamma0 N) 2)) =
      CuspForm.heckeULin 2 h := by
  rw [cuspHeckeGen, dif_pos h]; rfl

theorem cuspHeckeGen_of_not_dvd (ℓ : Nat.Primes) (h : ¬ (ℓ : ℕ) ∣ N) :
    (cuspHeckeGen N ℓ : Module.End ℂ (CuspForm (CongruenceSubgroup.Gamma0 N) 2)) =
      CuspForm.heckeTLin 2 ℓ.prop h := by
  rw [cuspHeckeGen, dif_neg h]; rfl

def cuspHeckeAeval : HeckeAlg →ₐ[ℤ] CuspForm.heckeAlgebra N 2 (∅ : Set ℕ) :=
  MvPolynomial.aeval (cuspHeckeGen N)

@[simp] theorem cuspHeckeAeval_heckeGen (ℓ : Nat.Primes) :
    cuspHeckeAeval N (heckeGen ℓ) = cuspHeckeGen N ℓ :=
  MvPolynomial.aeval_X _ ℓ

theorem range_cuspHeckeGen : Set.range (cuspHeckeGen N) =
    {x : CuspForm.heckeAlgebra N 2 (∅ : Set ℕ) |
      (x : Module.End ℂ (CuspForm (CongruenceSubgroup.Gamma0 N) 2)) ∈
        CuspForm.heckeGenerators N 2 (∅ : Set ℕ)} := by
  ext x
  constructor
  · rintro ⟨ℓ, rfl⟩
    by_cases h : (ℓ : ℕ) ∣ N
    · rw [Set.mem_setOf_eq, cuspHeckeGen_of_dvd N ℓ h]
      exact CuspForm.heckeULin_mem_heckeGenerators ℓ.prop h (Set.notMem_empty _)
    · rw [Set.mem_setOf_eq, cuspHeckeGen_of_not_dvd N ℓ h]
      exact CuspForm.heckeTLin_mem_heckeGenerators ℓ.prop h (Set.notMem_empty _)
  · rintro (⟨ℓ, hℓ, hℓN, -, hx⟩ | ⟨q, hqN, hq, -, hx⟩)
    · refine ⟨⟨ℓ, hℓ⟩, Subtype.ext ?_⟩
      rw [hx]; exact cuspHeckeGen_of_not_dvd N ⟨ℓ, hℓ⟩ hℓN
    · refine ⟨⟨q, hq⟩, Subtype.ext ?_⟩
      rw [hx]; exact cuspHeckeGen_of_dvd N ⟨q, hq⟩ hqN

theorem cuspHeckeAeval_range : (cuspHeckeAeval N).range = ⊤ := by
  rw [cuspHeckeAeval, ← Algebra.adjoin_range_eq_range_aeval, range_cuspHeckeGen]
  set H := CuspForm.heckeAlgebra N 2 (∅ : Set ℕ)
  apply Subalgebra.map_injective (f := H.val) Subtype.val_injective
  rw [AlgHom.map_adjoin, Algebra.map_top, Subalgebra.range_val]
  have himg : (H.val : H → Module.End ℂ (CuspForm (CongruenceSubgroup.Gamma0 N) 2)) ''
      {x : H | (x : Module.End ℂ _) ∈ CuspForm.heckeGenerators N 2 (∅ : Set ℕ)} =
      CuspForm.heckeGenerators N 2 (∅ : Set ℕ) := by
    ext T
    constructor
    · rintro ⟨x, hx, rfl⟩; exact hx
    · intro hT
      exact ⟨⟨T, Algebra.subset_adjoin hT⟩, hT, rfl⟩
  rw [himg]
  rfl

def cuspHeckeRep : HeckeAlg →+* Module.End ℂ (CuspForm (CongruenceSubgroup.Gamma0 N) 2) :=
  ((CuspForm.heckeAlgebra N 2 (∅ : Set ℕ)).val.comp (cuspHeckeAeval N)).toRingHom

theorem cuspHeckeRep_apply (t : HeckeAlg) :
    cuspHeckeRep N t = (cuspHeckeAeval N t : Module.End ℂ (CuspForm (CongruenceSubgroup.Gamma0 N) 2)) :=
  rfl

theorem cuspHeckeRep_heckeGen (ℓ : Nat.Primes) :
    cuspHeckeRep N (heckeGen ℓ) = (cuspHeckeGen N ℓ : Module.End ℂ _) := by
  rw [cuspHeckeRep_apply, cuspHeckeAeval_heckeGen]

theorem cuspHeckeRep_mem (t : HeckeAlg) : cuspHeckeRep N t ∈ CuspForm.heckeAlgebra N 2 (∅ : Set ℕ) :=
  (cuspHeckeAeval N t).2

theorem range_cuspHeckeRep :
    Set.range (cuspHeckeRep N) = (CuspForm.heckeAlgebra N 2 (∅ : Set ℕ) : Set _) := by
  ext T
  constructor
  · rintro ⟨t, rfl⟩; exact cuspHeckeRep_mem N t
  · intro hT
    have : (⟨T, hT⟩ : CuspForm.heckeAlgebra N 2 (∅ : Set ℕ)) ∈ (cuspHeckeAeval N).range := by
      rw [cuspHeckeAeval_range]; trivial
    obtain ⟨t, ht⟩ := this
    exact ⟨t, congrArg Subtype.val ht⟩

theorem cuspHeckeRep_comm (s t : HeckeAlg) :
    cuspHeckeRep N s * cuspHeckeRep N t = cuspHeckeRep N t * cuspHeckeRep N s := by
  rw [← map_mul, ← map_mul, mul_comm]

def dualHeckeRep : HeckeAlg →+* Module.End ℂ (Module.Dual ℂ (CuspForm (CongruenceSubgroup.Gamma0 N) 2)) where
  toFun t := (cuspHeckeRep N t).dualMap
  map_one' := by rw [map_one]; exact LinearMap.dualMap_id
  map_mul' s t := by
    rw [mul_comm s t, map_mul]
    exact (LinearMap.dualMap_comp_dualMap (cuspHeckeRep N s) (cuspHeckeRep N t)).symm
  map_zero' := by
    rw [map_zero]; exact map_zero (Module.Dual.transpose (R := ℂ))
  map_add' s t := by
    rw [map_add]; exact map_add (Module.Dual.transpose (R := ℂ)) _ _

theorem dualHeckeRep_apply (t : HeckeAlg) :
    dualHeckeRep N t = (cuspHeckeRep N t).dualMap := rfl

@[simp] theorem dualHeckeRep_apply_apply (t : HeckeAlg)
    (φ : Module.Dual ℂ (CuspForm (CongruenceSubgroup.Gamma0 N) 2))
    (f : CuspForm (CongruenceSubgroup.Gamma0 N) 2) :
    dualHeckeRep N t φ f = φ (cuspHeckeRep N t f) := rfl

theorem dualHeckeRep_C (a : ℤ) (φ : Module.Dual ℂ (CuspForm (CongruenceSubgroup.Gamma0 N) 2)) :
    dualHeckeRep N (MvPolynomial.C a) φ = a • φ := by
  ext f
  rw [dualHeckeRep_apply_apply, ← MvPolynomial.algebraMap_eq, eq_intCast, map_intCast,
    Module.End.intCast_apply, map_zsmul, LinearMap.smul_apply]

def PeriodLatticeHeckeStable : Prop :=
  ∀ (ℓ : Nat.Primes) (x : Module.Dual ℂ (CuspForm (CongruenceSubgroup.Gamma0 N) 2)),
    x ∈ periodLattice N → dualHeckeRep N (heckeGen ℓ) x ∈ periodLattice N

variable {N}

theorem periodLattice_stable_of (h : PeriodLatticeHeckeStable N) (t : HeckeAlg)
    (x : Module.Dual ℂ (CuspForm (CongruenceSubgroup.Gamma0 N) 2)) (hx : x ∈ periodLattice N) :
    dualHeckeRep N t x ∈ periodLattice N := by
  induction t using MvPolynomial.induction_on generalizing x with
  | C a => rw [dualHeckeRep_C]; exact Submodule.smul_mem _ a hx
  | add p q hp hq => rw [map_add, LinearMap.add_apply]; exact Submodule.add_mem _ (hp x hx) (hq x hx)
  | mul_X p ℓ hp =>
    rw [map_mul, Module.End.mul_apply]
    exact hp _ (h ℓ x hx)

variable (N)

def periodLatticeRestrict (h : PeriodLatticeHeckeStable N) (t : HeckeAlg) :
    Module.End ℤ (periodLattice N) :=
  ((dualHeckeRep N t).toAddMonoidHom.toIntLinearMap).restrict
    (fun x hx => periodLattice_stable_of h t x hx)

@[simp] theorem coe_periodLatticeRestrict_apply (h : PeriodLatticeHeckeStable N) (t : HeckeAlg)
    (x : periodLattice N) :
    ((periodLatticeRestrict N h t x : periodLattice N) : Module.Dual ℂ _) = dualHeckeRep N t x :=
  rfl

def periodLatticeRestrictHom (h : PeriodLatticeHeckeStable N) :
    HeckeAlg →+* Module.End ℤ (periodLattice N) where
  toFun := periodLatticeRestrict N h
  map_one' := LinearMap.ext fun x => Subtype.ext (by simp)
  map_mul' s t := LinearMap.ext fun x => Subtype.ext (by simp)
  map_zero' := LinearMap.ext fun x => Subtype.ext (by simp)
  map_add' s t := LinearMap.ext fun x => Subtype.ext (by simp)

open Classical in

def periodLatticeHeckeEnd : HeckeAlg →+* Module.End ℤ (periodLattice N) :=
  if h : PeriodLatticeHeckeStable N then periodLatticeRestrictHom N h
  else (Int.castRingHom (Module.End ℤ (periodLattice N))).comp
    (MvPolynomial.eval₂Hom (Int.castRingHom ℤ) (0 : Nat.Primes → ℤ))

variable {N}

theorem coe_periodLatticeHeckeEnd_apply (h : PeriodLatticeHeckeStable N) (t : HeckeAlg)
    (x : periodLattice N) :
    ((periodLatticeHeckeEnd N t x : periodLattice N) : Module.Dual ℂ _) = dualHeckeRep N t x := by
  have e : periodLatticeHeckeEnd N = periodLatticeRestrictHom N h := dif_pos h
  rw [e]; rfl

theorem periodLatticeHeckeEnd_of_not (h : ¬ PeriodLatticeHeckeStable N) (t : HeckeAlg)
    (x : periodLattice N) :
    periodLatticeHeckeEnd N t x = MvPolynomial.constantCoeff t • x := by
  have e : periodLatticeHeckeEnd N = (Int.castRingHom (Module.End ℤ (periodLattice N))).comp
      (MvPolynomial.eval₂Hom (Int.castRingHom ℤ) (0 : Nat.Primes → ℤ)) := dif_neg h
  rw [e, RingHom.comp_apply, MvPolynomial.eval₂Hom_zero_apply]
  simp only [eq_intCast, Int.cast_id, Module.End.intCast_apply]

variable (N)

@[implicit_reducible]
def periodLatticeModule : Module HeckeAlg (periodLattice N) :=
  Module.compHom (periodLattice N) (periodLatticeHeckeEnd N)

set_option synthInstance.maxHeartbeats 200000 in
theorem periodLatticeModule_smul_def (t : HeckeAlg) (x : periodLattice N) :
    (letI : Module HeckeAlg (periodLattice N) := periodLatticeModule N; t • x) =
      periodLatticeHeckeEnd N t x :=
  rfl

end Hecke

end ModularCurve

end