import Mathlib.RingTheory.MvPolynomial.Basic
import Mathlib.RingTheory.RootsOfUnity.PrimitiveRoots
import Mathlib.FieldTheory.Finite.GaloisField

set_option autoImplicit false

universe u

noncomputable section

open MvPolynomial

namespace DrinfeldCurve

variable (q : ℕ) (k : Type u) [Field k]

def drinfeldPoly : MvPolynomial (Fin 2) k := X 0 * X 1 ^ q - X 0 ^ q * X 1

def drinfeldIdeal : Ideal (MvPolynomial (Fin 2) k) := Ideal.span {drinfeldPoly q k - 1}

def CoordRing : Type u := MvPolynomial (Fin 2) k ⧸ drinfeldIdeal q k

instance instCommRingCoordRing : CommRing (CoordRing q k) :=
  inferInstanceAs (CommRing (MvPolynomial (Fin 2) k ⧸ drinfeldIdeal q k))

instance instAlgebraCoordRing : Algebra k (CoordRing q k) :=
  inferInstanceAs (Algebra k (MvPolynomial (Fin 2) k ⧸ drinfeldIdeal q k))

def mk : MvPolynomial (Fin 2) k →ₐ[k] CoordRing q k := Ideal.Quotient.mkₐ k (drinfeldIdeal q k)

theorem mk_surjective : Function.Surjective (mk q k) := Ideal.Quotient.mkₐ_surjective k _

theorem algHom_ext {A : Type*} [Semiring A] [Algebra k A] {f g : CoordRing q k →ₐ[k] A}
    (h : f.comp (mk q k) = g.comp (mk q k)) : f = g :=
  Ideal.Quotient.algHom_ext k h

theorem mk_drinfeldPoly : mk q k (drinfeldPoly q k) = 1 := by
  have h : mk q k (drinfeldPoly q k - 1) = 0 :=
    Ideal.Quotient.eq_zero_iff_mem.mpr (Ideal.subset_span (Set.mem_singleton _))
  rwa [map_sub, map_one, sub_eq_zero] at h

def x : CoordRing q k := mk q k (X 0)

def y : CoordRing q k := mk q k (X 1)

theorem relation : x q k * y q k ^ q - x q k ^ q * y q k = 1 := by
  have h := mk_drinfeldPoly q k
  rwa [drinfeldPoly, map_sub, map_mul, map_mul, map_pow, map_pow] at h

private theorem mk_mul_drinfeldPoly_sub_one (p : MvPolynomial (Fin 2) k) :
    mk q k (p * (drinfeldPoly q k - 1)) = 0 := by
  rw [map_mul, map_sub, map_one, mk_drinfeldPoly, sub_self, mul_zero]

def scalePoly (c : k) : MvPolynomial (Fin 2) k →ₐ[k] MvPolynomial (Fin 2) k :=
  aeval fun i => C c * X i

theorem scalePoly_X (c : k) (i : Fin 2) : scalePoly k c (X i) = C c * X i := by
  simp [scalePoly]

theorem scalePoly_drinfeldPoly (c : k) :
    scalePoly k c (drinfeldPoly q k) = C (c ^ (q + 1)) * drinfeldPoly q k := by
  simp only [drinfeldPoly, map_sub, map_mul, map_pow, scalePoly_X, mul_pow]
  ring

theorem scalePoly_one : scalePoly k 1 = AlgHom.id k _ := by
  apply MvPolynomial.algHom_ext
  intro i
  simp [scalePoly_X]

theorem scalePoly_mul (c d : k) : scalePoly k (c * d) = (scalePoly k c).comp (scalePoly k d) := by
  apply MvPolynomial.algHom_ext
  intro i
  simp only [scalePoly_X, AlgHom.comp_apply, map_mul, MvPolynomial.algHom_C, MvPolynomial.algebraMap_eq]
  ring

section Actions

variable [Fact q.Prime] [Algebra (GaloisField q 2) k]

def ofZMod : ZMod q →+* k :=
  (algebraMap (GaloisField q 2) k).comp (algebraMap (ZMod q) (GaloisField q 2))

theorem ofZMod_pow_card (a : ZMod q) : ofZMod q k a ^ q = ofZMod q k a := by
  rw [← map_pow, ZMod.pow_card]

private theorem charP_base : CharP k q :=
  charP_of_injective_algebraMap (algebraMap (GaloisField q 2) k).injective q

private theorem C_ofZMod_pow (a : ZMod q) :
    (C (ofZMod q k a) : MvPolynomial (Fin 2) k) ^ q = C (ofZMod q k a) := by
  rw [← map_pow, ofZMod_pow_card]

def substPoly (g : Matrix (Fin 2) (Fin 2) (ZMod q)) :
    MvPolynomial (Fin 2) k →ₐ[k] MvPolynomial (Fin 2) k :=
  aeval fun j => ∑ i, C (ofZMod q k (g i j)) * X i

theorem substPoly_X (g : Matrix (Fin 2) (Fin 2) (ZMod q)) (j : Fin 2) :
    substPoly q k g (X j) = ∑ i, C (ofZMod q k (g i j)) * X i := by
  simp [substPoly]

theorem substPoly_mul (g h : Matrix (Fin 2) (Fin 2) (ZMod q)) :
    substPoly q k (g * h) = (substPoly q k g).comp (substPoly q k h) := by
  apply MvPolynomial.algHom_ext
  intro j
  simp only [substPoly_X, AlgHom.comp_apply, map_mul, MvPolynomial.algHom_C, MvPolynomial.algebraMap_eq,
    Matrix.mul_apply, Fin.sum_univ_two, map_add]
  ring

theorem substPoly_one : substPoly q k 1 = AlgHom.id k _ := by
  apply MvPolynomial.algHom_ext
  intro j
  fin_cases j <;> simp [substPoly_X, Matrix.one_apply]

theorem substPoly_drinfeldPoly (g : Matrix (Fin 2) (Fin 2) (ZMod q)) :
    substPoly q k g (drinfeldPoly q k) = C (ofZMod q k g.det) * drinfeldPoly q k := by
  haveI := charP_base q k
  simp only [drinfeldPoly, map_sub, map_mul, map_pow, substPoly_X, Fin.sum_univ_two, Matrix.det_fin_two]
  rw [add_pow_char, add_pow_char]
  simp only [mul_pow, C_ofZMod_pow]
  ring

private def slEnd (g : Matrix.SpecialLinearGroup (Fin 2) (ZMod q)) : CoordRing q k →ₐ[k] CoordRing q k :=
  Ideal.Quotient.liftₐ (drinfeldIdeal q k) ((mk q k).comp (substPoly q k (g : Matrix (Fin 2) (Fin 2) (ZMod q))))
    (by
      intro a ha
      rw [drinfeldIdeal, Ideal.mem_span_singleton'] at ha
      obtain ⟨b, rfl⟩ := ha
      rw [AlgHom.comp_apply, map_mul, map_sub, map_one, substPoly_drinfeldPoly,
        Matrix.SpecialLinearGroup.det_coe, map_one, C_1, one_mul]
      exact mk_mul_drinfeldPoly_sub_one q k _)

private theorem slEnd_mk (g : Matrix.SpecialLinearGroup (Fin 2) (ZMod q)) (p : MvPolynomial (Fin 2) k) :
    slEnd q k g (mk q k p) = mk q k (substPoly q k (g : Matrix (Fin 2) (Fin 2) (ZMod q)) p) :=
  rfl

private def slEndHom : Matrix.SpecialLinearGroup (Fin 2) (ZMod q) →* (CoordRing q k →ₐ[k] CoordRing q k) where
  toFun := slEnd q k
  map_one' := by
    apply algHom_ext
    apply MvPolynomial.algHom_ext
    intro j
    simp [slEnd_mk, Matrix.SpecialLinearGroup.coe_one, substPoly_one]
  map_mul' g h := by
    apply algHom_ext
    apply MvPolynomial.algHom_ext
    intro j
    simp [slEnd_mk, Matrix.SpecialLinearGroup.coe_mul, substPoly_mul]

def slAction : Matrix.SpecialLinearGroup (Fin 2) (ZMod q) →* (CoordRing q k ≃ₐ[k] CoordRing q k) :=
  (AlgEquiv.algHomUnitsEquiv k (CoordRing q k)).toMonoidHom.comp (slEndHom q k).toHomUnits

theorem slAction_mk (g : Matrix.SpecialLinearGroup (Fin 2) (ZMod q)) (p : MvPolynomial (Fin 2) k) :
    slAction q k g (mk q k p) = mk q k (substPoly q k (g : Matrix (Fin 2) (Fin 2) (ZMod q)) p) :=
  slEnd_mk q k g p

theorem slAction_x (g : Matrix.SpecialLinearGroup (Fin 2) (ZMod q)) :
    slAction q k g (x q k) =
      algebraMap k _ (ofZMod q k ((g : Matrix (Fin 2) (Fin 2) (ZMod q)) 0 0)) * x q k +
        algebraMap k _ (ofZMod q k ((g : Matrix (Fin 2) (Fin 2) (ZMod q)) 1 0)) * y q k := by
  simp [x, y, slAction_mk, substPoly_X, Fin.sum_univ_two]

theorem slAction_y (g : Matrix.SpecialLinearGroup (Fin 2) (ZMod q)) :
    slAction q k g (y q k) =
      algebraMap k _ (ofZMod q k ((g : Matrix (Fin 2) (Fin 2) (ZMod q)) 0 1)) * x q k +
        algebraMap k _ (ofZMod q k ((g : Matrix (Fin 2) (Fin 2) (ZMod q)) 1 1)) * y q k := by
  simp [x, y, slAction_mk, substPoly_X, Fin.sum_univ_two]

theorem card_galoisField_two : Nat.card (GaloisField q 2) = q ^ 2 := GaloisField.card q 2 two_ne_zero

theorem natCard_rootsOfUnity_galoisField_two :
    Nat.card (rootsOfUnity (q + 1) (GaloisField q 2)) = q + 1 := by
  have hq0 : q ≠ 0 := (Fact.out : q.Prime).ne_zero
  have hsplit : (q - 1) * (q + 1) + 1 = q ^ 2 := by
    obtain ⟨m, rfl⟩ := Nat.exists_eq_add_one_of_ne_zero hq0
    rw [Nat.add_sub_cancel]
    ring
  have hcard : Nat.card (GaloisField q 2)ˣ = (q - 1) * (q + 1) := by
    rw [Nat.card_units, card_galoisField_two, ← hsplit, Nat.add_sub_cancel]
  obtain ⟨g, hg⟩ := IsCyclic.exists_generator (α := (GaloisField q 2)ˣ)
  have hord : orderOf g = (q - 1) * (q + 1) := by
    rw [← hcard, orderOf_eq_card_of_forall_mem_zpowers hg]
  have hprim : IsPrimitiveRoot (g ^ (q - 1)) (q + 1) :=
    (IsPrimitiveRoot.orderOf g).pow (hord ▸ orderOf_pos g) hord
  exact hprim.card_rootsOfUnity'

def scalarOf (ζ : (GaloisField q 2)ˣ) : k := algebraMap (GaloisField q 2) k ζ

theorem scalarOf_pow_of_mem {ζ : (GaloisField q 2)ˣ} (hζ : ζ ∈ rootsOfUnity (q + 1) (GaloisField q 2)) :
    scalarOf q k ζ ^ (q + 1) = 1 := by
  rw [mem_rootsOfUnity] at hζ
  rw [scalarOf, ← map_pow, ← Units.val_pow_eq_pow_val, hζ, Units.val_one, map_one]

private theorem scalarOf_one : scalarOf q k 1 = 1 := by
  simp [scalarOf]

private theorem scalarOf_mul (ζ η : (GaloisField q 2)ˣ) :
    scalarOf q k (ζ * η) = scalarOf q k ζ * scalarOf q k η := by
  simp [scalarOf]

private def muEnd (ζ : rootsOfUnity (q + 1) (GaloisField q 2)) : CoordRing q k →ₐ[k] CoordRing q k :=
  Ideal.Quotient.liftₐ (drinfeldIdeal q k) ((mk q k).comp (scalePoly k (scalarOf q k ζ)))
    (by
      intro a ha
      rw [drinfeldIdeal, Ideal.mem_span_singleton'] at ha
      obtain ⟨b, rfl⟩ := ha
      rw [AlgHom.comp_apply, map_mul, map_sub, map_one, scalePoly_drinfeldPoly, scalarOf_pow_of_mem q k ζ.2,
        C_1, one_mul]
      exact mk_mul_drinfeldPoly_sub_one q k _)

private theorem muEnd_mk (ζ : rootsOfUnity (q + 1) (GaloisField q 2)) (p : MvPolynomial (Fin 2) k) :
    muEnd q k ζ (mk q k p) = mk q k (scalePoly k (scalarOf q k ζ) p) :=
  rfl

private def muEndHom : rootsOfUnity (q + 1) (GaloisField q 2) →* (CoordRing q k →ₐ[k] CoordRing q k) where
  toFun := muEnd q k
  map_one' := by
    apply algHom_ext
    apply MvPolynomial.algHom_ext
    intro i
    simp [muEnd_mk, scalarOf_one, scalePoly_one]
  map_mul' ζ η := by
    apply algHom_ext
    apply MvPolynomial.algHom_ext
    intro i
    simp [muEnd_mk, scalarOf_mul, scalePoly_mul]

def muAction : rootsOfUnity (q + 1) (GaloisField q 2) →* (CoordRing q k ≃ₐ[k] CoordRing q k) :=
  (AlgEquiv.algHomUnitsEquiv k (CoordRing q k)).toMonoidHom.comp (muEndHom q k).toHomUnits

theorem muAction_mk (ζ : rootsOfUnity (q + 1) (GaloisField q 2)) (p : MvPolynomial (Fin 2) k) :
    muAction q k ζ (mk q k p) = mk q k (scalePoly k (scalarOf q k ζ) p) :=
  muEnd_mk q k ζ p

theorem muAction_x (ζ : rootsOfUnity (q + 1) (GaloisField q 2)) :
    muAction q k ζ (x q k) = algebraMap k _ (scalarOf q k ζ) * x q k := by
  simp [x, muAction_mk, scalePoly_X]

theorem muAction_y (ζ : rootsOfUnity (q + 1) (GaloisField q 2)) :
    muAction q k ζ (y q k) = algebraMap k _ (scalarOf q k ζ) * y q k := by
  simp [y, muAction_mk, scalePoly_X]

theorem substPoly_comp_scalePoly (g : Matrix (Fin 2) (Fin 2) (ZMod q)) (c : k) :
    (substPoly q k g).comp (scalePoly k c) = (scalePoly k c).comp (substPoly q k g) := by
  apply MvPolynomial.algHom_ext
  intro j
  simp only [AlgHom.comp_apply, scalePoly_X, substPoly_X, map_mul, MvPolynomial.algHom_C,
    MvPolynomial.algebraMap_eq, Fin.sum_univ_two, map_add]
  ring

theorem slAction_comm_muAction (g : Matrix.SpecialLinearGroup (Fin 2) (ZMod q))
    (ζ : rootsOfUnity (q + 1) (GaloisField q 2)) :
    slAction q k g * muAction q k ζ = muAction q k ζ * slAction q k g := by
  ext z
  obtain ⟨p, rfl⟩ := mk_surjective q k z
  simp only [AlgEquiv.mul_apply, slAction_mk, muAction_mk]
  exact congrArg (mk q k) (AlgHom.congr_fun (substPoly_comp_scalePoly q k (g : Matrix (Fin 2) (Fin 2) (ZMod q)) (scalarOf q k ζ)) p)

theorem scalePoly_substPoly_drinfeldPoly (g : Matrix (Fin 2) (Fin 2) (ZMod q)) (c : k) :
    scalePoly k c (substPoly q k g (drinfeldPoly q k)) =
      C (ofZMod q k g.det * c ^ (q + 1)) * drinfeldPoly q k := by
  simp only [substPoly_drinfeldPoly, map_mul, scalePoly_drinfeldPoly, MvPolynomial.algHom_C,
    MvPolynomial.algebraMap_eq]
  ring

def hChar : Matrix.GeneralLinearGroup (Fin 2) (ZMod q) × (GaloisField q 2)ˣ →* (GaloisField q 2)ˣ :=
  ((Units.map (algebraMap (ZMod q) (GaloisField q 2)).toMonoidHom).comp
      (Matrix.GeneralLinearGroup.det.comp (MonoidHom.fst _ _))) *
    ((powMonoidHom (q + 1)).comp (MonoidHom.snd _ _))

theorem coe_hChar_apply (p : Matrix.GeneralLinearGroup (Fin 2) (ZMod q) × (GaloisField q 2)ˣ) :
    ((hChar q p : (GaloisField q 2)ˣ) : GaloisField q 2) =
      algebraMap (ZMod q) (GaloisField q 2) (p.1 : Matrix (Fin 2) (Fin 2) (ZMod q)).det *
        ((p.2 : (GaloisField q 2)ˣ) : GaloisField q 2) ^ (q + 1) := by
  simp [hChar, Matrix.GeneralLinearGroup.val_det_apply]

def hSubgroup : Subgroup (Matrix.GeneralLinearGroup (Fin 2) (ZMod q) × (GaloisField q 2)ˣ) := (hChar q).ker

theorem mem_hSubgroup_iff (p : Matrix.GeneralLinearGroup (Fin 2) (ZMod q) × (GaloisField q 2)ˣ) :
    p ∈ hSubgroup q ↔ hChar q p = 1 :=
  MonoidHom.mem_ker

private theorem hFactor {p : Matrix.GeneralLinearGroup (Fin 2) (ZMod q) × (GaloisField q 2)ˣ}
    (hp : p ∈ hSubgroup q) :
    ofZMod q k (p.1 : Matrix (Fin 2) (Fin 2) (ZMod q)).det * scalarOf q k p.2 ^ (q + 1) = 1 := by
  have h := congrArg (fun u : (GaloisField q 2)ˣ => algebraMap (GaloisField q 2) k u) ((mem_hSubgroup_iff q p).mp hp)
  simp only [coe_hChar_apply, Units.val_one, map_one, map_mul, map_pow] at h
  simpa [ofZMod, scalarOf] using h

private def hEnd (p : hSubgroup q) : CoordRing q k →ₐ[k] CoordRing q k :=
  Ideal.Quotient.liftₐ (drinfeldIdeal q k)
    ((mk q k).comp ((scalePoly k (scalarOf q k (p : Matrix.GeneralLinearGroup (Fin 2) (ZMod q) ×
      (GaloisField q 2)ˣ).2)).comp (substPoly q k ((p : Matrix.GeneralLinearGroup (Fin 2) (ZMod q) ×
      (GaloisField q 2)ˣ).1 : Matrix (Fin 2) (Fin 2) (ZMod q)))))
    (by
      intro a ha
      rw [drinfeldIdeal, Ideal.mem_span_singleton'] at ha
      obtain ⟨b, rfl⟩ := ha
      rw [AlgHom.comp_apply, AlgHom.comp_apply, map_mul, map_sub, map_one, map_mul, map_sub, map_one,
        scalePoly_substPoly_drinfeldPoly, hFactor q k p.2, C_1, one_mul]
      exact mk_mul_drinfeldPoly_sub_one q k _)

private theorem hEnd_mk (p : hSubgroup q) (f : MvPolynomial (Fin 2) k) :
    hEnd q k p (mk q k f) =
      mk q k (scalePoly k (scalarOf q k (p : Matrix.GeneralLinearGroup (Fin 2) (ZMod q) × (GaloisField q 2)ˣ).2)
        (substPoly q k ((p : Matrix.GeneralLinearGroup (Fin 2) (ZMod q) × (GaloisField q 2)ˣ).1 :
          Matrix (Fin 2) (Fin 2) (ZMod q)) f)) :=
  rfl

private def hEndHom : hSubgroup q →* (CoordRing q k →ₐ[k] CoordRing q k) where
  toFun := hEnd q k
  map_one' := by
    apply algHom_ext
    apply MvPolynomial.algHom_ext
    intro i
    simp [hEnd_mk, scalarOf_one, scalePoly_one, substPoly_one]
  map_mul' p₁ p₂ := by
    apply algHom_ext
    apply MvPolynomial.algHom_ext
    intro j
    simp only [AlgHom.comp_apply, AlgHom.mul_apply, hEnd_mk]
    simp only [Subgroup.coe_mul, Prod.fst_mul, Prod.snd_mul, Units.val_mul, scalarOf_mul, substPoly_X, scalePoly_X,
      map_mul, map_add, MvPolynomial.algHom_C, MvPolynomial.algebraMap_eq, Matrix.mul_apply, Fin.sum_univ_two]
    ring

def hAction : hSubgroup q →* (CoordRing q k ≃ₐ[k] CoordRing q k) :=
  (AlgEquiv.algHomUnitsEquiv k (CoordRing q k)).toMonoidHom.comp (hEndHom q k).toHomUnits

theorem hAction_mk (p : hSubgroup q) (f : MvPolynomial (Fin 2) k) :
    hAction q k p (mk q k f) =
      mk q k (scalePoly k (scalarOf q k (p : Matrix.GeneralLinearGroup (Fin 2) (ZMod q) × (GaloisField q 2)ˣ).2)
        (substPoly q k ((p : Matrix.GeneralLinearGroup (Fin 2) (ZMod q) × (GaloisField q 2)ˣ).1 :
          Matrix (Fin 2) (Fin 2) (ZMod q)) f)) :=
  hEnd_mk q k p f

theorem toGL_one_mem_hSubgroup (g : Matrix.SpecialLinearGroup (Fin 2) (ZMod q)) :
    (Matrix.SpecialLinearGroup.toGL g, (1 : (GaloisField q 2)ˣ)) ∈ hSubgroup q := by
  rw [mem_hSubgroup_iff, Units.ext_iff, coe_hChar_apply]
  simp [Matrix.SpecialLinearGroup.coe_GL_coe_matrix]

theorem one_mem_hSubgroup_of_mem (ζ : rootsOfUnity (q + 1) (GaloisField q 2)) :
    ((1 : Matrix.GeneralLinearGroup (Fin 2) (ZMod q)), (ζ : (GaloisField q 2)ˣ)) ∈ hSubgroup q := by
  rw [mem_hSubgroup_iff, Units.ext_iff, coe_hChar_apply]
  have hζ := (mem_rootsOfUnity _ _).mp ζ.2
  simp [← Units.val_pow_eq_pow_val, hζ]

theorem slAction_eq_hAction (g : Matrix.SpecialLinearGroup (Fin 2) (ZMod q)) :
    slAction q k g = hAction q k ⟨_, toGL_one_mem_hSubgroup q g⟩ := by
  ext z
  obtain ⟨f, rfl⟩ := mk_surjective q k z
  rw [slAction_mk, hAction_mk]
  simp [scalarOf_one, scalePoly_one, Matrix.SpecialLinearGroup.coe_GL_coe_matrix]

theorem muAction_eq_hAction (ζ : rootsOfUnity (q + 1) (GaloisField q 2)) :
    muAction q k ζ = hAction q k ⟨_, one_mem_hSubgroup_of_mem q ζ⟩ := by
  ext z
  obtain ⟨f, rfl⟩ := mk_surjective q k z
  rw [muAction_mk, hAction_mk]
  simp [substPoly_one]

end Actions

end DrinfeldCurve

end