import Mathlib
import Definitions.Def_CohCarrier_Inst
import Definitions.Def_CuspForm_HeckeOperatorFormsGammaH
import Definitions.Def_ModularForm_AtkinLehnerDatum
import Definitions.Def_FLTPrelim_Modularity

set_option autoImplicit false

noncomputable section

open scoped MatrixGroups ModularForm
open CongruenceSubgroup CohCarrier

namespace CuspForm

section HeckeRing

variable {M : ℕ} [NeZero M] {H : Subgroup (ZMod M)ˣ}

def heckeGenH (S : Set ℕ) (k : ℤ) :
    Gen M S → (CuspForm (GammaH M H) k →ₗ[ℂ] CuspForm (GammaH M H) k)
  | .T _ hℓ _ hℓM => heckeTLinH k hℓ hℓM
  | .U q _ _ => heckeULinH k q
  | .dia d => diamondLinH k d

@[simp] theorem heckeGenH_T (S : Set ℕ) (k : ℤ) (ℓ : ℕ) (hℓ : ℓ.Prime) (hℓS : ℓ ∉ S) (hℓM : ¬ ℓ ∣ M) :
    heckeGenH (H := H) S k (.T ℓ hℓ hℓS hℓM) = heckeTLinH k hℓ hℓM := rfl

@[simp] theorem heckeGenH_U (S : Set ℕ) (k : ℤ) (q : ℕ) (hq : q.Prime) (hqM : q ∣ M) :
    heckeGenH (H := H) S k (.U q hq hqM) = heckeULinH k q := rfl

@[simp] theorem heckeGenH_dia (S : Set ℕ) (k : ℤ) (d : (ZMod M)ˣ) :
    heckeGenH (H := H) S k (.dia d) = diamondLinH k d := rfl

variable (M H) in

def heckeRingH (k : ℤ) : Subring (CuspForm (GammaH M H) k →ₗ[ℂ] CuspForm (GammaH M H) k) :=
  Subring.closure (Set.range (heckeGenH (H := H) (∅ : Set ℕ) k))

theorem heckeGenH_mem_heckeRingH (S : Set ℕ) (k : ℤ) (g : Gen M S) :
    heckeGenH (H := H) S k g ∈ heckeRingH M H k := by
  refine Subring.subset_closure ?_
  cases g with
  | T ℓ hℓ hℓS hℓM => exact ⟨.T ℓ hℓ (Set.notMem_empty ℓ) hℓM, rfl⟩
  | U q hq hqM => exact ⟨.U q hq hqM, rfl⟩
  | dia d => exact ⟨.dia d, rfl⟩

end HeckeRing

section Lattice

variable (M : ℕ) [NeZero M] (H : Subgroup (ZMod M)ˣ) (k : ℤ) (p : ℕ) (A : Subring ℂ)

def twoCuspIntegralSet : Set (CuspForm (GammaH M H) k) :=
  {f | ∀ t ∈ heckeRingH M H k, ∀ (W : ModularForm.AtkinLehnerDatum M p) (n : ℕ),
    ModularFormClass.qCoeff (⇑(t f)) n ∈ A ∧
      ModularFormClass.qCoeff (ModularForm.alSlash W k ⇑(t f)) n ∈ A}

variable {M H k p A} in
theorem mem_twoCuspIntegralSet_iff (f : CuspForm (GammaH M H) k) :
    f ∈ twoCuspIntegralSet M H k p A ↔
      ∀ t ∈ heckeRingH M H k, ∀ (W : ModularForm.AtkinLehnerDatum M p) (n : ℕ),
        ModularFormClass.qCoeff (⇑(t f)) n ∈ A ∧
          ModularFormClass.qCoeff (ModularForm.alSlash W k ⇑(t f)) n ∈ A :=
  Iff.rfl

theorem zero_mem_twoCuspIntegralSet : (0 : CuspForm (GammaH M H) k) ∈ twoCuspIntegralSet M H k p A := by
  intro t _ W n
  have h0 : ModularFormClass.qCoeff (0 : UpperHalfPlane → ℂ) n = 0 := by
    simp [ModularFormClass.qCoeff, UpperHalfPlane.qExpansion_zero]
  have hW : ModularForm.alSlash W k (0 : UpperHalfPlane → ℂ) = 0 := by
    rw [ModularForm.alSlash_def]; exact SlashAction.zero_slash k _
  refine ⟨?_, ?_⟩
  · rw [map_zero, CuspForm.coe_zero, h0]; exact A.zero_mem
  · rw [map_zero, CuspForm.coe_zero, hW, h0]; exact A.zero_mem

variable {M H k p A} in

theorem heckeRingH_apply_mem_twoCuspIntegralSet {t : CuspForm (GammaH M H) k →ₗ[ℂ] CuspForm (GammaH M H) k}
    (ht : t ∈ heckeRingH M H k) {f : CuspForm (GammaH M H) k} (hf : f ∈ twoCuspIntegralSet M H k p A) :
    t f ∈ twoCuspIntegralSet M H k p A := by
  intro t' ht' W n
  have hmul : t' * t ∈ heckeRingH M H k := (heckeRingH M H k).mul_mem ht' ht
  exact hf (t' * t) hmul W n

def twoCuspLattice : Submodule A (CuspForm (GammaH M H) k) :=
  Submodule.span A (twoCuspIntegralSet M H k p A)

theorem twoCuspIntegralSet_subset_twoCuspLattice :
    twoCuspIntegralSet M H k p A ⊆ (twoCuspLattice M H k p A : Set (CuspForm (GammaH M H) k)) :=
  Submodule.subset_span

variable {M H k p A} in

theorem heckeRingH_apply_mem_twoCuspLattice {t : CuspForm (GammaH M H) k →ₗ[ℂ] CuspForm (GammaH M H) k}
    (ht : t ∈ heckeRingH M H k) {f : CuspForm (GammaH M H) k} (hf : f ∈ twoCuspLattice M H k p A) :
    t f ∈ twoCuspLattice M H k p A := by
  induction hf using Submodule.span_induction with
  | mem x hx => exact Submodule.subset_span (heckeRingH_apply_mem_twoCuspIntegralSet ht hx)
  | zero => rw [map_zero]; exact Submodule.zero_mem _
  | add x y _ _ hx hy => rw [map_add]; exact Submodule.add_mem _ hx hy
  | smul a x _ hx =>
      rw [show t (a • x) = a • t x from by
        rw [Subring.smul_def, Subring.smul_def, LinearMap.map_smul]]
      exact Submodule.smul_mem _ a hx

variable {M H k p A} in

def twoCuspEnd (t : heckeRingH M H k) : twoCuspLattice M H k p A →ₗ[A] twoCuspLattice M H k p A where
  toFun x := ⟨(t : CuspForm (GammaH M H) k →ₗ[ℂ] _) x, heckeRingH_apply_mem_twoCuspLattice t.2 x.2⟩
  map_add' x y := Subtype.ext (by simp)
  map_smul' a x := Subtype.ext (by
    simp only [Submodule.coe_smul, RingHom.id_apply, Subring.smul_def, LinearMap.map_smul])

@[simp] theorem coe_twoCuspEnd_apply (t : heckeRingH M H k) (x : twoCuspLattice M H k p A) :
    ((twoCuspEnd t x : twoCuspLattice M H k p A) : CuspForm (GammaH M H) k) =
      (t : CuspForm (GammaH M H) k →ₗ[ℂ] _) x :=
  rfl

theorem twoCuspEnd_one : twoCuspEnd (p := p) (A := A) (1 : heckeRingH M H k) = LinearMap.id :=
  LinearMap.ext fun _ => Subtype.ext rfl

theorem twoCuspEnd_mul (t t' : heckeRingH M H k) :
    twoCuspEnd (p := p) (A := A) (t * t') = twoCuspEnd t ∘ₗ twoCuspEnd t' :=
  LinearMap.ext fun _ => Subtype.ext rfl

end Lattice

section Reduction

variable (M : ℕ) [NeZero M] (H : Subgroup (ZMod M)ˣ) (k : ℤ) (p : ℕ) (A : Subring ℂ) (I : Ideal A)

def TwoCuspForms : Type :=
  twoCuspLattice M H k p A ⧸ (I • ⊤ : Submodule A (twoCuspLattice M H k p A))

instance instAddCommGroupTwoCuspForms : AddCommGroup (TwoCuspForms M H k p A I) :=
  inferInstanceAs (AddCommGroup (twoCuspLattice M H k p A ⧸ (I • ⊤ : Submodule A (twoCuspLattice M H k p A))))

instance instModuleTwoCuspForms : Module A (TwoCuspForms M H k p A I) :=
  inferInstanceAs (Module A (twoCuspLattice M H k p A ⧸ (I • ⊤ : Submodule A (twoCuspLattice M H k p A))))

instance instModuleQuotientTwoCuspForms : Module (A ⧸ I) (TwoCuspForms M H k p A I) :=
  (Module.isTorsionBySet_quotient_ideal_smul (twoCuspLattice M H k p A) I).module

instance instIsScalarTowerSelfTwoCuspForms : IsScalarTower A A (TwoCuspForms M H k p A I) :=
  inferInstanceAs
    (IsScalarTower A A (twoCuspLattice M H k p A ⧸ (I • ⊤ : Submodule A (twoCuspLattice M H k p A))))

instance instIsScalarTowerTwoCuspForms : IsScalarTower A (A ⧸ I) (TwoCuspForms M H k p A I) :=
  ⟨fun a b x => by
    obtain ⟨c, rfl⟩ := Ideal.Quotient.mk_surjective b
    change (a * c) • x = a • (c • x)
    exact mul_smul a c x⟩

variable {M H k p A} in

def twoCuspReduce : twoCuspLattice M H k p A →ₗ[A] TwoCuspForms M H k p A I :=
  (I • ⊤ : Submodule A (twoCuspLattice M H k p A)).mkQ

theorem twoCuspReduce_surjective : Function.Surjective (twoCuspReduce (M := M) (H := H) (k := k) (p := p) (A := A) I) :=
  Submodule.mkQ_surjective _

theorem ker_twoCuspReduce :
    LinearMap.ker (twoCuspReduce (M := M) (H := H) (k := k) (p := p) (A := A) I) =
      (I • ⊤ : Submodule A (twoCuspLattice M H k p A)) :=
  Submodule.ker_mkQ _

theorem twoCuspReduce_eq_zero_iff (x : twoCuspLattice M H k p A) :
    twoCuspReduce I x = 0 ↔ x ∈ (I • ⊤ : Submodule A (twoCuspLattice M H k p A)) :=
  Submodule.Quotient.mk_eq_zero _

theorem mk_smul_twoCuspReduce (a : A) (x : twoCuspLattice M H k p A) :
    (Ideal.Quotient.mk I a) • twoCuspReduce I x = twoCuspReduce I (a • x) :=
  rfl

variable {M H k p A} in

def twoCuspEndModAux (t : heckeRingH M H k) : TwoCuspForms M H k p A I →ₗ[A] TwoCuspForms M H k p A I :=
  Submodule.mapQ (I • ⊤ : Submodule A (twoCuspLattice M H k p A)) (I • ⊤ : Submodule A (twoCuspLattice M H k p A))
    (twoCuspEnd t) (by
      rw [← Submodule.map_le_iff_le_comap, Submodule.map_smul'']
      exact Submodule.smul_mono le_rfl le_top)

variable {M H k p A} in

def twoCuspEndMod (t : heckeRingH M H k) : TwoCuspForms M H k p A I →ₗ[A ⧸ I] TwoCuspForms M H k p A I :=
  LinearMap.extendScalarsOfSurjective Ideal.Quotient.mk_surjective (twoCuspEndModAux I t)

@[simp] theorem twoCuspEndMod_reduce (t : heckeRingH M H k) (x : twoCuspLattice M H k p A) :
    twoCuspEndMod I t (twoCuspReduce I x) = twoCuspReduce I (twoCuspEnd t x) :=
  rfl

theorem twoCuspEndMod_comp_reduce (t : heckeRingH M H k) :
    (twoCuspEndMod (p := p) (A := A) I t).restrictScalars A ∘ₗ twoCuspReduce I =
      twoCuspReduce I ∘ₗ twoCuspEnd t :=
  rfl

variable {M H k p A} in

def twoCuspGenMod (S : Set ℕ) (g : Gen M S) : TwoCuspForms M H k p A I →ₗ[A ⧸ I] TwoCuspForms M H k p A I :=
  twoCuspEndMod I ⟨heckeGenH S k g, heckeGenH_mem_heckeRingH S k g⟩

theorem twoCuspGenMod_reduce (S : Set ℕ) (g : Gen M S) (x : twoCuspLattice M H k p A) :
    twoCuspGenMod I S g (twoCuspReduce I x) =
      twoCuspReduce I (twoCuspEnd ⟨heckeGenH S k g, heckeGenH_mem_heckeRingH S k g⟩ x) :=
  rfl

variable {M H k p A} in

def twoCuspEigenspace (S : Set ℕ) (χ : Gen M S → A ⧸ I) : Submodule (A ⧸ I) (TwoCuspForms M H k p A I) where
  carrier := {ω | ∀ g : Gen M S, twoCuspGenMod I S g ω = χ g • ω}
  zero_mem' g := by rw [map_zero, smul_zero]
  add_mem' {x y} hx hy g := by rw [map_add, hx g, hy g, smul_add]
  smul_mem' c {x} hx g := by rw [LinearMap.map_smul, hx g, smul_comm]

variable {M H k p A} in
theorem mem_twoCuspEigenspace_iff (S : Set ℕ) (χ : Gen M S → A ⧸ I) (ω : TwoCuspForms M H k p A I) :
    ω ∈ twoCuspEigenspace I S χ ↔ ∀ g : Gen M S, twoCuspGenMod I S g ω = χ g • ω :=
  Iff.rfl

end Reduction

end CuspForm

end