import Mathlib
import Definitions.Def_MvFormalGroup_NegV2
import Definitions.Def_MvFormalGroup_CartierModule
import Definitions.Def_MvFormalGroup_CartierModuleHomothety
import Definitions.Def_CerednikDrinfeld_SpecialFormalModule

set_option autoImplicit false

noncomputable section

universe u

namespace CerednikDrinfeld

namespace FormalODModule

open MvFormalGroup MvFormalGroup.CartierModule

variable {p : ℕ} [hp : Fact p.Prime] {B : Type u} [CommRing B]

def gradedPiece (X : FormalODModule p B) (j : Zp2 p →+* B) (n : ℕ) :
    AddSubgroup (CartierModule p X.F) where
  carrier := {f | ∀ c : GaloisField p 2,
    endAct (X.actEnd (WittVector.teichmuller p c)) f =
      homothety (j (WittVector.teichmuller p c) ^ p ^ n) f}
  zero_mem' := fun c => by simp only [map_zero]
  add_mem' := fun {f g} hf hg c => by simp only [map_add, hf c, hg c]
  neg_mem' := fun {f} hf c => by simp only [map_neg, hf c]

theorem mem_gradedPiece_iff (X : FormalODModule p B) (j : Zp2 p →+* B) (n : ℕ)
    (f : CartierModule p X.F) :
    f ∈ X.gradedPiece j n ↔ ∀ c : GaloisField p 2,
      endAct (X.actEnd (WittVector.teichmuller p c)) f =
        homothety (j (WittVector.teichmuller p c) ^ p ^ n) f :=
  Iff.rfl

theorem pow_card_galoisField (c : GaloisField p 2) : c ^ p ^ 2 = c := by
  letI : Fintype (GaloisField p 2) := Fintype.ofFinite _
  have hcard : Fintype.card (GaloisField p 2) = p ^ 2 := by
    rw [Fintype.card_eq_nat_card, GaloisField.card p 2 two_ne_zero]
  rw [← hcard]
  exact FiniteField.pow_card c

theorem apply_teichmuller_pow_pow_add_two (j : Zp2 p →+* B) (c : GaloisField p 2) (n : ℕ) :
    j (WittVector.teichmuller p c) ^ p ^ (n + 2) = j (WittVector.teichmuller p c) ^ p ^ n := by
  rw [pow_add, pow_mul, ← map_pow, ← map_pow, ← map_pow, ← map_pow, ← pow_mul, mul_comm, pow_mul,
    pow_card_galoisField]

theorem gradedPiece_add_two (X : FormalODModule p B) (j : Zp2 p →+* B) (n : ℕ) :
    X.gradedPiece j (n + 2) = X.gradedPiece j n := by
  ext f
  simp only [mem_gradedPiece_iff, apply_teichmuller_pow_pow_add_two]

theorem gradedPiece_add_two_mul (X : FormalODModule p B) (j : Zp2 p →+* B) (n k : ℕ) :
    X.gradedPiece j (n + 2 * k) = X.gradedPiece j n := by
  induction k with
  | zero => rw [mul_zero, add_zero]
  | succ k ih => rw [mul_add, mul_one, ← add_assoc, gradedPiece_add_two, ih]

end FormalODModule

end CerednikDrinfeld

end