import Mathlib

set_option autoImplicit false

open scoped TensorProduct
open Coalgebra Bialgebra WithConv

universe u v w

namespace PDivisibleGroup.Hopf

section NSMul

variable (R : Type u) [CommSemiring R] (A : Type v) [CommSemiring A] [Bialgebra R A]

noncomputable def nsmulAlgHom (n : ℕ) : A →ₐ[R] A :=
  ((toConv (AlgHom.id R A)) ^ n).ofConv

theorem toConv_nsmulAlgHom (n : ℕ) :
    toConv (nsmulAlgHom R A n) = toConv (AlgHom.id R A) ^ n := rfl

theorem nsmulAlgHom_zero :
    nsmulAlgHom R A 0 = (Algebra.ofId R A).comp (counitAlgHom R A) := by
  change ((toConv (AlgHom.id R A)) ^ 0).ofConv = _
  rw [pow_zero]
  rfl

@[simp] theorem nsmulAlgHom_zero_apply (a : A) :
    nsmulAlgHom R A 0 a = algebraMap R A (counit (R := R) a) := by
  rw [nsmulAlgHom_zero]
  rfl

@[simp] theorem nsmulAlgHom_one : nsmulAlgHom R A 1 = AlgHom.id R A := by
  change ((toConv (AlgHom.id R A)) ^ 1).ofConv = _
  rw [pow_one]

theorem nsmulAlgHom_succ (n : ℕ) :
    nsmulAlgHom R A (n + 1) = (toConv (nsmulAlgHom R A n) * toConv (AlgHom.id R A)).ofConv := by
  change ((toConv (AlgHom.id R A)) ^ (n + 1)).ofConv = _
  rw [pow_succ]
  rfl

theorem nsmulAlgHom_add (m n : ℕ) :
    nsmulAlgHom R A (m + n) = (toConv (nsmulAlgHom R A m) * toConv (nsmulAlgHom R A n)).ofConv := by
  change ((toConv (AlgHom.id R A)) ^ (m + n)).ofConv = _
  rw [pow_add]
  rfl

theorem nsmulAlgHom_mul (m n : ℕ) :
    nsmulAlgHom R A (m * n) = ((toConv (nsmulAlgHom R A m)) ^ n).ofConv := by
  change ((toConv (AlgHom.id R A)) ^ (m * n)).ofConv = _
  rw [pow_mul]
  rfl

theorem toLinearMap_nsmulAlgHom (n : ℕ) :
    (nsmulAlgHom R A n).toLinearMap = ((toConv (LinearMap.id : A →ₗ[R] A)) ^ n).ofConv := by
  have h := AlgHom.toLinearMap_convPow (toConv (AlgHom.id R A)) n
  exact congrArg WithConv.ofConv h

theorem toConv_toLinearMap_nsmulAlgHom (n : ℕ) :
    toConv (nsmulAlgHom R A n).toLinearMap = (toConv (LinearMap.id : A →ₗ[R] A)) ^ n := by
  rw [toLinearMap_nsmulAlgHom]

variable {R A}

theorem algHom_comp_convPow {B : Type w} [CommSemiring B] [Algebra R B] (g : A →ₐ[R] B)
    (f : WithConv (A →ₐ[R] A)) (n : ℕ) :
    g.comp (f ^ n).ofConv = ((toConv (g.comp f.ofConv)) ^ n).ofConv := by
  induction n with
  | zero =>
    rw [pow_zero, pow_zero]
    ext a
    change g (algebraMap R A (counit (R := R) a)) = algebraMap R B (counit (R := R) a)
    exact g.commutes _
  | succ n ih =>
    rw [pow_succ, pow_succ, AlgHom.comp_convMul_distrib, ih]

theorem convPow_comp_bialgHom {B : Type w} [CommSemiring B] [Bialgebra R B]
    {C : Type*} [CommSemiring C] [Algebra R C]
    (f : WithConv (A →ₐ[R] C)) (g : B →ₐc[R] A) (n : ℕ) :
    (f ^ n).ofConv.comp (g : B →ₐ[R] A) = ((toConv (f.ofConv.comp (g : B →ₐ[R] A))) ^ n).ofConv := by
  induction n with
  | zero =>
    rw [pow_zero, pow_zero]
    ext b
    change algebraMap R C (counit (R := R) (g b)) = algebraMap R C (counit (R := R) b)
    rw [CoalgHomClass.counit_comp_apply]
  | succ n ih =>
    rw [pow_succ, pow_succ, AlgHom.convMul_comp_bialgHom_distrib, ih]

theorem toConv_pow_eq_toConv_comp_nsmulAlgHom {L : Type w} [CommSemiring L] [Algebra R L]
    (f : A →ₐ[R] L) (n : ℕ) :
    (toConv f) ^ n = toConv (f.comp (nsmulAlgHom R A n)) := by
  apply WithConv.ofConv_injective
  change ((toConv f) ^ n).ofConv = f.comp (((toConv (AlgHom.id R A)) ^ n).ofConv)
  rw [algHom_comp_convPow]
  rfl

theorem pow_apply_eq_apply_nsmulAlgHom {L : Type w} [CommSemiring L] [Algebra R L]
    (f : WithConv (A →ₐ[R] L)) (n : ℕ) (a : A) :
    (f ^ n) a = f (nsmulAlgHom R A n a) := by
  have h := toConv_pow_eq_toConv_comp_nsmulAlgHom f.ofConv n
  rw [toConv_ofConv] at h
  rw [h]
  rfl

theorem nsmulAlgHom_comp_bialgHom {B : Type w} [CommSemiring B] [Bialgebra R B]
    (g : B →ₐc[R] A) (n : ℕ) :
    (nsmulAlgHom R A n).comp (g : B →ₐ[R] A) = (g : B →ₐ[R] A).comp (nsmulAlgHom R B n) := by
  change ((toConv (AlgHom.id R A)) ^ n).ofConv.comp (g : B →ₐ[R] A) =
    (g : B →ₐ[R] A).comp (((toConv (AlgHom.id R B)) ^ n).ofConv)
  rw [convPow_comp_bialgHom, algHom_comp_convPow]
  rfl

theorem nsmulAlgHom_bialgHom_apply {B : Type w} [CommSemiring B] [Bialgebra R B]
    (g : B →ₐc[R] A) (n : ℕ) (b : B) :
    nsmulAlgHom R A n (g b) = g (nsmulAlgHom R B n b) :=
  DFunLike.congr_fun (nsmulAlgHom_comp_bialgHom g n) b

@[simp] theorem counit_nsmulAlgHom (n : ℕ) (a : A) :
    counit (R := R) (nsmulAlgHom R A n a) = counit (R := R) a := by
  have h : (counitAlgHom R A).comp (nsmulAlgHom R A n) = counitAlgHom R A := by
    change (counitAlgHom R A).comp (((toConv (AlgHom.id R A)) ^ n).ofConv) = _
    rw [algHom_comp_convPow]
    have h1 : toConv ((counitAlgHom R A).comp (AlgHom.id R A)) = (1 : WithConv (A →ₐ[R] R)) := by
      rw [AlgHom.convOne_def]
      exact congrArg _ (AlgHom.ext fun a => rfl)
    rw [h1, one_pow]
    ext a
    change algebraMap R R (counit (R := R) a) = counit (R := R) a
    simp
  exact DFunLike.congr_fun h a

theorem nsmulAlgHom_algebraMap (n : ℕ) (r : R) :
    nsmulAlgHom R A n (algebraMap R A r) = algebraMap R A r :=
  (nsmulAlgHom R A n).commutes r

end NSMul

section Ideals

variable (R : Type u) [CommRing R] (A : Type v) [CommRing A] [Bialgebra R A]

def augIdeal : Ideal A := RingHom.ker (counitAlgHom R A)

theorem mem_augIdeal_iff (a : A) : a ∈ augIdeal R A ↔ counit (R := R) a = 0 := Iff.rfl

theorem sub_algebraMap_counit_mem_augIdeal (a : A) :
    a - algebraMap R A (counit (R := R) a) ∈ augIdeal R A := by
  rw [mem_augIdeal_iff, map_sub, counit_algebraMap, sub_self]

theorem nsmulAlgHom_mem_augIdeal {a : A} (ha : a ∈ augIdeal R A) (n : ℕ) :
    nsmulAlgHom R A n a ∈ augIdeal R A := by
  rw [mem_augIdeal_iff] at ha ⊢
  rw [counit_nsmulAlgHom, ha]

noncomputable def torsionIdeal (n : ℕ) : Ideal A := (augIdeal R A).map (nsmulAlgHom R A n)

theorem torsionIdeal_eq_span (n : ℕ) :
    torsionIdeal R A n = Ideal.span (nsmulAlgHom R A n '' (augIdeal R A : Set A)) := rfl

theorem nsmulAlgHom_mem_torsionIdeal {a : A} (ha : a ∈ augIdeal R A) (n : ℕ) :
    nsmulAlgHom R A n a ∈ torsionIdeal R A n :=
  Ideal.mem_map_of_mem _ ha

theorem torsionIdeal_le_augIdeal (n : ℕ) : torsionIdeal R A n ≤ augIdeal R A := by
  rw [torsionIdeal, Ideal.map_le_iff_le_comap]
  intro a ha
  exact nsmulAlgHom_mem_augIdeal R A ha n

@[simp] theorem torsionIdeal_one : torsionIdeal R A 1 = augIdeal R A := by
  rw [torsionIdeal, nsmulAlgHom_one]
  exact Ideal.map_id _

@[simp] theorem torsionIdeal_zero : torsionIdeal R A 0 = ⊥ := by
  rw [← le_bot_iff, torsionIdeal, Ideal.map_le_iff_le_comap]
  intro a ha
  change nsmulAlgHom R A 0 a ∈ (⊥ : Ideal A)
  rw [nsmulAlgHom_zero_apply, (mem_augIdeal_iff R A a).1 ha, map_zero]
  exact Ideal.zero_mem _

theorem apply_eq_zero_of_mem_torsionIdeal_of_pow_eq_one {L : Type w} [CommRing L] [Algebra R L]
    (f : WithConv (A →ₐ[R] L)) (n : ℕ) (hf : f ^ n = 1) {a : A} (ha : a ∈ torsionIdeal R A n) :
    f a = 0 := by
  have hvan : ∀ b ∈ augIdeal R A, f (nsmulAlgHom R A n b) = 0 := by
    intro b hb
    rw [← pow_apply_eq_apply_nsmulAlgHom, hf, AlgHom.convOne_apply, (mem_augIdeal_iff R A b).1 hb,
      map_zero]
  have hle : torsionIdeal R A n ≤ RingHom.ker (f.ofConv : A →ₐ[R] L) := by
    rw [torsionIdeal, Ideal.map_le_iff_le_comap]
    intro b hb
    exact hvan b hb
  exact hle ha

end Ideals

end PDivisibleGroup.Hopf

structure PDivisibleGroup (R : Type) [CommRing R] (p h : ℕ) : Type 1 where

  level : ℕ → Type
  [instCommRing : ∀ v, CommRing (level v)]
  [instHopfAlgebra : ∀ v, HopfAlgebra R (level v)]
  [instIsCocomm : ∀ v, Coalgebra.IsCocomm R (level v)]
  [instModuleFinite : ∀ v, Module.Finite R (level v)]
  [instModuleFree : ∀ v, Module.Free R (level v)]

  transition : ∀ v, level (v + 1) →ₐc[R] level v
  transition_surjective : ∀ v, Function.Surjective (transition v)

  finrank_level : ∀ v, Module.finrank R (level v) = p ^ (v * h)

  ker_transition : ∀ v,
    RingHom.ker (transition v) = PDivisibleGroup.Hopf.torsionIdeal R (level (v + 1)) (p ^ v)

attribute [instance] PDivisibleGroup.instCommRing PDivisibleGroup.instHopfAlgebra
  PDivisibleGroup.instIsCocomm PDivisibleGroup.instModuleFinite PDivisibleGroup.instModuleFree

namespace PDivisibleGroup

variable {R : Type} [CommRing R] {p h : ℕ} (G : PDivisibleGroup R p h)

def height (_G : PDivisibleGroup R p h) : ℕ := h

@[simp] theorem height_eq : G.height = h := rfl

instance instModuleFlat (v : ℕ) : Module.Flat R (G.level v) := inferInstance

abbrev transitionAlgHom (v : ℕ) : G.level (v + 1) →ₐ[R] G.level v :=
  (G.transition v : G.level (v + 1) →ₐ[R] G.level v)

theorem transitionAlgHom_apply (v : ℕ) (a : G.level (v + 1)) :
    G.transitionAlgHom v a = G.transition v a := rfl

theorem transitionAlgHom_surjective (v : ℕ) : Function.Surjective (G.transitionAlgHom v) :=
  G.transition_surjective v

@[simp] theorem counit_transition (v : ℕ) (a : G.level (v + 1)) :
    counit (R := R) (G.transition v a) = counit (R := R) a :=
  CoalgHomClass.counit_comp_apply _ a

theorem finrank_level_zero : Module.finrank R (G.level 0) = 1 := by
  rw [G.finrank_level 0, zero_mul, pow_zero]

theorem mem_ker_transition_iff (v : ℕ) (a : G.level (v + 1)) :
    G.transition v a = 0 ↔ a ∈ Hopf.torsionIdeal R (G.level (v + 1)) (p ^ v) := by
  rw [← G.ker_transition v, RingHom.mem_ker]

theorem transition_nsmulAlgHom_eq_zero (v : ℕ) {a : G.level (v + 1)}
    (ha : a ∈ Hopf.augIdeal R (G.level (v + 1))) :
    G.transition v (Hopf.nsmulAlgHom R (G.level (v + 1)) (p ^ v) a) = 0 :=
  (G.mem_ker_transition_iff v _).2 (Hopf.nsmulAlgHom_mem_torsionIdeal R _ ha _)

theorem nsmulAlgHom_transition (v n : ℕ) (a : G.level (v + 1)) :
    Hopf.nsmulAlgHom R (G.level v) n (G.transition v a) =
      G.transition v (Hopf.nsmulAlgHom R (G.level (v + 1)) n a) :=
  Hopf.nsmulAlgHom_bialgHom_apply (G.transition v) n a

theorem nsmulAlgHom_pow_level (v : ℕ) :
    Hopf.nsmulAlgHom R (G.level v) (p ^ v) =
      (Algebra.ofId R (G.level v)).comp (counitAlgHom R (G.level v)) := by
  refine AlgHom.ext fun b => ?_
  obtain ⟨a, rfl⟩ := G.transition_surjective v b
  have hdec : a = algebraMap R _ (counit (R := R) a) + (a - algebraMap R _ (counit (R := R) a)) := by
    abel
  rw [G.nsmulAlgHom_transition]
  conv_lhs => rw [hdec]
  rw [map_add, map_add, Hopf.nsmulAlgHom_algebraMap,
    G.transition_nsmulAlgHom_eq_zero v (Hopf.sub_algebraMap_counit_mem_augIdeal R _ a),
    add_zero, AlgHom.comp_apply, Algebra.ofId_apply, Bialgebra.counitAlgHom_apply, G.counit_transition]
  exact (G.transitionAlgHom v).commutes _

theorem nsmulAlgHom_pow_level_apply (v : ℕ) (b : G.level v) :
    Hopf.nsmulAlgHom R (G.level v) (p ^ v) b = algebraMap R (G.level v) (counit (R := R) b) := by
  rw [G.nsmulAlgHom_pow_level v]
  rfl

theorem toConv_linearMapId_pow_level (v : ℕ) :
    (toConv (LinearMap.id : G.level v →ₗ[R] G.level v)) ^ (p ^ v) = 1 := by
  rw [← Hopf.toConv_toLinearMap_nsmulAlgHom, G.nsmulAlgHom_pow_level v]
  rfl

theorem point_pow_eq_one (v : ℕ) {L : Type w} [CommRing L] [Algebra R L]
    (f : WithConv (G.level v →ₐ[R] L)) : f ^ (p ^ v) = 1 := by
  have h := Hopf.toConv_pow_eq_toConv_comp_nsmulAlgHom f.ofConv (p ^ v)
  rw [toConv_ofConv] at h
  rw [h, G.nsmulAlgHom_pow_level v, AlgHom.convOne_def, ← AlgHom.comp_assoc, Algebra.comp_ofId]

theorem comp_transition_injective (v : ℕ) {L : Type w} [CommRing L] [Algebra R L] :
    Function.Injective fun g : G.level v →ₐ[R] L => g.comp (G.transitionAlgHom v) := by
  intro g g' hgg'
  refine AlgHom.ext fun b => ?_
  obtain ⟨a, rfl⟩ := G.transition_surjective v b
  exact DFunLike.congr_fun hgg' a

theorem exists_comp_transition_eq (v : ℕ) {L : Type w} [CommRing L] [Algebra R L]
    (f : WithConv (G.level (v + 1) →ₐ[R] L)) (hf : f ^ (p ^ v) = 1) :
    ∃ g : G.level v →ₐ[R] L, g.comp (G.transitionAlgHom v) = f.ofConv := by
  have hle : RingHom.ker (G.transitionAlgHom v).toRingHom ≤ RingHom.ker f.ofConv.toRingHom := by
    intro a ha
    have ha' : a ∈ Hopf.torsionIdeal R (G.level (v + 1)) (p ^ v) := by
      rw [← G.ker_transition v]
      exact ha
    exact Hopf.apply_eq_zero_of_mem_torsionIdeal_of_pow_eq_one R _ f (p ^ v) hf ha'
  exact ⟨AlgHom.liftOfSurjective (G.transitionAlgHom v) (G.transition_surjective v) f.ofConv hle,
    AlgHom.liftOfSurjective_comp _ _ _ _⟩

theorem exists_comp_transition_eq_iff (v : ℕ) {L : Type w} [CommRing L] [Algebra R L]
    (f : WithConv (G.level (v + 1) →ₐ[R] L)) :
    (∃ g : G.level v →ₐ[R] L, g.comp (G.transitionAlgHom v) = f.ofConv) ↔ f ^ (p ^ v) = 1 := by
  refine ⟨?_, G.exists_comp_transition_eq v f⟩
  rintro ⟨g, hg⟩
  have hf : f = toConv (g.comp (G.transitionAlgHom v)) := by rw [hg, toConv_ofConv]
  rw [hf, ← WithConv.ofConv_injective.eq_iff]
  rw [show toConv (g.comp (G.transitionAlgHom v)) = toConv ((toConv g).ofConv.comp (G.transitionAlgHom v))
      from rfl, ← Hopf.convPow_comp_bialgHom, G.point_pow_eq_one v (toConv g)]
  rw [AlgHom.convOne_def, AlgHom.convOne_def]
  change ((Algebra.ofId R L).comp (counitAlgHom R (G.level v))).comp (G.transitionAlgHom v) = _
  refine AlgHom.ext fun a => ?_
  change algebraMap R L (counit (R := R) (G.transition v a)) = algebraMap R L (counit (R := R) a)
  rw [G.counit_transition]

end PDivisibleGroup