import Mathlib
import Definitions.Def_CerednikDrinfeld_FormalUpperHalfPlaneCharts

set_option autoImplicit false

noncomputable section

open Polynomial

namespace CerednikDrinfeld
namespace FormalOmega

section ChartRings

variable (R : Type) [CommRing R]

abbrev vertexDiscr (q : ℕ) : R[X] := X ^ q - X

abbrev chartVRing (q : ℕ) : Type := Localization.Away (vertexDiscr R q)

namespace chartVRing

variable (q : ℕ)

def ζ : chartVRing R q := algebraMap R[X] (chartVRing R q) X

theorem algebraMap_vertexDiscr : algebraMap R[X] (chartVRing R q) (vertexDiscr R q) = ζ R q ^ q - ζ R q := by
  show algebraMap R[X] (chartVRing R q) (X ^ q - X) = _
  rw [map_sub, map_pow]; rfl

theorem isUnit_ζ_pow_sub : IsUnit (ζ R q ^ q - ζ R q) := by
  rw [← algebraMap_vertexDiscr]
  exact IsLocalization.Away.algebraMap_isUnit (vertexDiscr R q)

def point : (chartV R q).obj (chartVRing R q) := ⟨ζ R q, isUnit_ζ_pow_sub R q⟩

variable {R q}

theorem isUnit_aeval_of_mem_powers {B : Type} [CommRing B] [Algebra R B] (x : (chartV R q).obj B)
    (y : Submonoid.powers (vertexDiscr R q)) : IsUnit (Polynomial.aeval (R := R) x.1 (y : R[X])) := by
  obtain ⟨_, n, rfl⟩ := y
  rw [map_pow]
  refine IsUnit.pow n ?_
  have : Polynomial.aeval (R := R) x.1 (vertexDiscr R q) = x.1 ^ q - x.1 := by
    show Polynomial.aeval (R := R) x.1 (X ^ q - X) = _
    rw [map_sub (Polynomial.aeval (R := R) x.1), map_pow (Polynomial.aeval (R := R) x.1), Polynomial.aeval_X]
  rw [this]; exact x.2

def lift {B : Type} [CommRing B] [Algebra R B] (x : (chartV R q).obj B) : chartVRing R q →ₐ[R] B :=
  IsLocalization.liftAlgHom (M := Submonoid.powers (vertexDiscr R q)) (f := Polynomial.aeval (R := R) x.1)
    (isUnit_aeval_of_mem_powers x)

@[simp] theorem lift_ζ {B : Type} [CommRing B] [Algebra R B] (x : (chartV R q).obj B) : lift x (ζ R q) = x.1 := by
  rw [lift, IsLocalization.liftAlgHom_apply, ζ, IsLocalization.lift_eq]
  exact Polynomial.aeval_X x.1

variable (R q) in

def corepEquiv (B : Type) [CommRing B] [Algebra R B] :
    (chartVRing R q →ₐ[R] B) ≃ (chartV R q).obj B where
  toFun φ := (chartV R q).map φ (point R q)
  invFun x := lift x
  left_inv φ := by
    apply IsLocalization.algHom_ext (Submonoid.powers (vertexDiscr R q))
    apply Polynomial.algHom_ext
    show lift ((chartV R q).map φ (point R q)) (algebraMap R[X] _ X) = φ (algebraMap R[X] _ X)
    exact lift_ζ _
  right_inv x := Subtype.ext (lift_ζ x)

@[simp] theorem corepEquiv_apply_val (B : Type) [CommRing B] [Algebra R B] (φ : chartVRing R q →ₐ[R] B) :
    (corepEquiv R q B φ).1 = φ (ζ R q) := rfl

theorem corepEquiv_natural {B : Type} [CommRing B] [Algebra R B] {B' : Type} [CommRing B'] [Algebra R B']
    (f : B →ₐ[R] B') (φ : chartVRing R q →ₐ[R] B) :
    corepEquiv R q B' ((AlgFunctor.corep (𝒪 := R) (chartVRing R q)).map f φ) = (chartV R q).map f (corepEquiv R q B φ) :=
  Subtype.ext rfl

variable (R q) in

def corepNatTrans : AlgFunctor.NatTrans (AlgFunctor.corep (𝒪 := R) (chartVRing R q)) (chartV R q) where
  app B _ _ φ := corepEquiv R q B φ
  naturality f φ := corepEquiv_natural f φ

end chartVRing

abbrev edgeRel (π : R) : MvPolynomial (Fin 2) R := MvPolynomial.X 0 * MvPolynomial.X 1 - MvPolynomial.C π

def edgeQuot (π : R) : Type := MvPolynomial (Fin 2) R ⧸ Ideal.span {edgeRel R π}

instance edgeQuot.instCommRing (π : R) : CommRing (edgeQuot R π) :=
  show CommRing (MvPolynomial (Fin 2) R ⧸ Ideal.span {edgeRel R π}) from inferInstance

instance edgeQuot.instAlgebra (π : R) : Algebra R (edgeQuot R π) :=
  show Algebra R (MvPolynomial (Fin 2) R ⧸ Ideal.span {edgeRel R π}) from inferInstance

def edgeQuot.mk (π : R) : MvPolynomial (Fin 2) R →ₐ[R] edgeQuot R π :=
  Ideal.Quotient.mkₐ R (Ideal.span {edgeRel R π})

theorem edgeQuot.mk_edgeRel (π : R) : edgeQuot.mk R π (edgeRel R π) = 0 :=
  Ideal.Quotient.eq_zero_iff_mem.mpr (Ideal.subset_span (Set.mem_singleton _))

abbrev edgeDiscr (q : ℕ) : MvPolynomial (Fin 2) R :=
  (MvPolynomial.X 0 ^ (q - 1) - 1) * (MvPolynomial.X 1 ^ (q - 1) - 1)

def edgeQuot.discr (π : R) (q : ℕ) : edgeQuot R π := edgeQuot.mk R π (edgeDiscr R q)

abbrev chartERing (π : R) (q : ℕ) : Type := Localization.Away (edgeQuot.discr R π q)

namespace chartERing

variable (π : R) (q : ℕ)

def ofPoly : MvPolynomial (Fin 2) R →ₐ[R] chartERing R π q :=
  (Algebra.algHom R (edgeQuot R π) (chartERing R π q)).comp (edgeQuot.mk R π)

def ξ : chartERing R π q := ofPoly R π q (MvPolynomial.X 0)

def η : chartERing R π q := ofPoly R π q (MvPolynomial.X 1)

theorem ofPoly_edgeRel : ofPoly R π q (edgeRel R π) = 0 := by
  rw [ofPoly, AlgHom.comp_apply, edgeQuot.mk_edgeRel, map_zero]

theorem ξ_mul_η : ξ R π q * η R π q = algebraMap R (chartERing R π q) π := by
  have hC : ofPoly R π q (MvPolynomial.C π) = algebraMap R (chartERing R π q) π := (ofPoly R π q).commutes π
  have h2 : ofPoly R π q (edgeRel R π) = ξ R π q * η R π q - algebraMap R (chartERing R π q) π := by
    show ofPoly R π q (MvPolynomial.X 0 * MvPolynomial.X 1 - MvPolynomial.C π) = _
    rw [map_sub (ofPoly R π q), map_mul (ofPoly R π q), hC]; rfl
  rw [← sub_eq_zero, ← h2, ofPoly_edgeRel]

theorem ofPoly_edgeDiscr :
    ofPoly R π q (edgeDiscr R q) = (ξ R π q ^ (q - 1) - 1) * (η R π q ^ (q - 1) - 1) := by
  show ofPoly R π q ((MvPolynomial.X 0 ^ (q - 1) - 1) * (MvPolynomial.X 1 ^ (q - 1) - 1)) = _
  rw [map_mul (ofPoly R π q), map_sub (ofPoly R π q), map_sub (ofPoly R π q), map_pow (ofPoly R π q),
    map_pow (ofPoly R π q), map_one]; rfl

theorem isUnit_discr : IsUnit ((ξ R π q ^ (q - 1) - 1) * (η R π q ^ (q - 1) - 1)) := by
  rw [← ofPoly_edgeDiscr]
  exact IsLocalization.Away.algebraMap_isUnit (edgeQuot.discr R π q)

def point : (chartE R π q).obj (chartERing R π q) :=
  ⟨(ξ R π q, η R π q), ξ_mul_η R π q, isUnit_of_mul_isUnit_left (isUnit_discr R π q),
    isUnit_of_mul_isUnit_right (isUnit_discr R π q)⟩

variable {R π q}

abbrev evalPoly {B : Type} [CommRing B] [Algebra R B] (x : (chartE R π q).obj B) : MvPolynomial (Fin 2) R →ₐ[R] B :=
  MvPolynomial.aeval (fun i : Fin 2 => if i = 0 then x.1.1 else x.1.2)

@[simp] theorem evalPoly_X_zero {B : Type} [CommRing B] [Algebra R B] (x : (chartE R π q).obj B) :
    evalPoly x (MvPolynomial.X 0) = x.1.1 := by simp [evalPoly, MvPolynomial.aeval_X]

@[simp] theorem evalPoly_X_one {B : Type} [CommRing B] [Algebra R B] (x : (chartE R π q).obj B) :
    evalPoly x (MvPolynomial.X 1) = x.1.2 := by simp [evalPoly, MvPolynomial.aeval_X]

theorem evalPoly_edgeRel {B : Type} [CommRing B] [Algebra R B] (x : (chartE R π q).obj B) :
    evalPoly x (edgeRel R π) = 0 := by
  have hC : evalPoly x (MvPolynomial.C π) = algebraMap R B π := (evalPoly x).commutes π
  show evalPoly x (MvPolynomial.X 0 * MvPolynomial.X 1 - MvPolynomial.C π) = 0
  rw [map_sub (evalPoly x), map_mul (evalPoly x), evalPoly_X_zero, evalPoly_X_one, hC, x.2.1, sub_self]

def evalQuot {B : Type} [CommRing B] [Algebra R B] (x : (chartE R π q).obj B) : edgeQuot R π →ₐ[R] B :=
  Ideal.Quotient.liftₐ (Ideal.span {edgeRel R π}) (evalPoly x) (fun a ha => by
    obtain ⟨c, rfl⟩ := Ideal.mem_span_singleton'.mp ha
    rw [map_mul (evalPoly x), evalPoly_edgeRel, mul_zero])

theorem evalQuot_mk {B : Type} [CommRing B] [Algebra R B] (x : (chartE R π q).obj B) (p : MvPolynomial (Fin 2) R) :
    evalQuot x (edgeQuot.mk R π p) = evalPoly x p := rfl

theorem isUnit_evalQuot_of_mem_powers {B : Type} [CommRing B] [Algebra R B] (x : (chartE R π q).obj B)
    (y : Submonoid.powers (edgeQuot.discr R π q)) :
    IsUnit (evalQuot x (y : edgeQuot R π)) := by
  obtain ⟨_, n, rfl⟩ := y
  rw [map_pow]
  refine IsUnit.pow n ?_
  rw [edgeQuot.discr, evalQuot_mk]
  show IsUnit (evalPoly x ((MvPolynomial.X 0 ^ (q - 1) - 1) * (MvPolynomial.X 1 ^ (q - 1) - 1)))
  rw [map_mul (evalPoly x), map_sub (evalPoly x), map_sub (evalPoly x), map_pow (evalPoly x), map_pow (evalPoly x),
    map_one, evalPoly_X_zero, evalPoly_X_one]
  exact x.2.2.1.mul x.2.2.2

def lift {B : Type} [CommRing B] [Algebra R B] (x : (chartE R π q).obj B) : chartERing R π q →ₐ[R] B :=
  IsLocalization.liftAlgHom (M := Submonoid.powers (edgeQuot.discr R π q))
    (f := evalQuot x) (isUnit_evalQuot_of_mem_powers x)

theorem lift_ofPoly {B : Type} [CommRing B] [Algebra R B] (x : (chartE R π q).obj B) (p : MvPolynomial (Fin 2) R) :
    lift x (ofPoly R π q p) = evalPoly x p := by
  rw [lift, IsLocalization.liftAlgHom_apply, ofPoly, AlgHom.comp_apply]
  exact (IsLocalization.lift_eq _ _).trans (evalQuot_mk x p)

@[simp] theorem lift_ξ {B : Type} [CommRing B] [Algebra R B] (x : (chartE R π q).obj B) : lift x (ξ R π q) = x.1.1 := by
  rw [ξ, lift_ofPoly, evalPoly_X_zero]

@[simp] theorem lift_η {B : Type} [CommRing B] [Algebra R B] (x : (chartE R π q).obj B) : lift x (η R π q) = x.1.2 := by
  rw [η, lift_ofPoly, evalPoly_X_one]

variable (R π q) in

def corepEquiv (B : Type) [CommRing B] [Algebra R B] :
    (chartERing R π q →ₐ[R] B) ≃ (chartE R π q).obj B where
  toFun φ := (chartE R π q).map φ (point R π q)
  invFun x := lift x
  left_inv φ := by
    apply IsLocalization.algHom_ext (Submonoid.powers (edgeQuot.discr R π q))
    apply Ideal.Quotient.algHom_ext
    apply MvPolynomial.algHom_ext
    intro i
    show lift ((chartE R π q).map φ (point R π q)) (ofPoly R π q (MvPolynomial.X i)) = φ (ofPoly R π q (MvPolynomial.X i))
    fin_cases i
    · exact lift_ξ _
    · exact lift_η _
  right_inv x := Subtype.ext (Prod.ext (lift_ξ x) (lift_η x))

@[simp] theorem corepEquiv_apply_val (B : Type) [CommRing B] [Algebra R B] (φ : chartERing R π q →ₐ[R] B) :
    (corepEquiv R π q B φ).1 = (φ (ξ R π q), φ (η R π q)) := rfl

theorem corepEquiv_natural {B : Type} [CommRing B] [Algebra R B] {B' : Type} [CommRing B'] [Algebra R B']
    (f : B →ₐ[R] B') (φ : chartERing R π q →ₐ[R] B) :
    corepEquiv R π q B' ((AlgFunctor.corep (𝒪 := R) (chartERing R π q)).map f φ) = (chartE R π q).map f (corepEquiv R π q B φ) :=
  Subtype.ext rfl

variable (R π q) in

def corepNatTrans : AlgFunctor.NatTrans (AlgFunctor.corep (𝒪 := R) (chartERing R π q)) (chartE R π q) where
  app B _ _ φ := corepEquiv R π q B φ
  naturality f φ := corepEquiv_natural f φ

end chartERing

section Level

variable (𝒪 : Type) [CommRing 𝒪] (π : 𝒪) (q n : ℕ)

abbrev truncRing : Type := 𝒪 ⧸ Ideal.span {π ^ (n + 1)}

abbrev chartVRingMod : Type := chartVRing (truncRing 𝒪 π n) q

abbrev chartERingMod : Type :=
  chartERing (truncRing 𝒪 π n) (Ideal.Quotient.mk (Ideal.span {π ^ (n + 1)}) π) q

theorem isNilpotent_algebraMap_of_isScalarTower (B : Type) [CommRing B] [Algebra 𝒪 B]
    [Algebra (truncRing 𝒪 π n) B] [IsScalarTower 𝒪 (truncRing 𝒪 π n) B] : IsNilpotent (algebraMap 𝒪 B π) := by
  refine ⟨n + 1, ?_⟩
  rw [← map_pow, IsScalarTower.algebraMap_apply 𝒪 (truncRing 𝒪 π n) B,
    Ideal.Quotient.algebraMap_eq, Ideal.Quotient.eq_zero_iff_mem.mpr (Ideal.subset_span (Set.mem_singleton _)),
    map_zero]

end Level

end ChartRings

end FormalOmega
end CerednikDrinfeld

end