import Mathlib

namespace CategoryTheory

universe v u

section Definition

structure GradedMonad (C : Type u) [Category.{v} C] where
  toFunctor : ℕ → C ⥤ C
  η : 𝟭 _ ⟶ toFunctor 0
  μ (n k : ℕ) : toFunctor k ⋙ toFunctor n ⟶ toFunctor (n + k)
  assoc : ∀ (X : C) (n k m : ℕ),
      (toFunctor n).map (NatTrans.app (μ k m) X) ≫ (μ n (k + m)).app X =
      (μ n k).app ((toFunctor m).obj X) ≫ (μ (n + k) m).app X ≫ eqToHom (by rw [Nat.add_assoc]) :=
    by cat_disch
  left_unit : ∀ (X : C) (n : ℕ),
      η.app ((toFunctor n).obj X) ≫ (μ 0 n).app X =
      𝟙 ((toFunctor n).obj X) ≫ eqToHom (by rw [Nat.zero_add]) :=
    by cat_disch
  right_unit : ∀ (X : C) (n : ℕ),
      (toFunctor n).map (η.app X) ≫ (μ n 0).app X =
      𝟙 ((toFunctor n).obj X) :=
    by cat_disch

namespace GradedMonad

variable {C : Type u} [Category.{v} C]

instance : CoeFun (GradedMonad C) (fun _ => ℕ → C ⥤ C) where
  coe M := M.toFunctor

@[reassoc]
lemma unit_naturality (M : GradedMonad C) {X Y : C} (f : X ⟶ Y) :
    f ≫ M.η.app Y = M.η.app X ≫ (M 0).map f :=
  M.η.naturality _

@[reassoc]
lemma multiplication_naturality (M : GradedMonad C) {X Y : C} (f : X ⟶ Y) {n k : ℕ} :
    (M n).map ((M k).map f) ≫ (M.μ n k).app Y = (M.μ n k).app X ≫ (M (n + k)).map f :=
  (M.μ n k).naturality _

end GradedMonad

end Definition

section Instances

namespace GradedMonad

variable {C : Type u} [Category.{v} C]

def ofMonad (M : Monad C) : GradedMonad C where
  toFunctor _ := M.toFunctor
  η := M.η
  μ _ _ := M.μ
  assoc X _ _ _ := by simp only [Functor.comp_obj, eqToHom_refl, Category.comp_id]; exact M.assoc X
  left_unit X _ := by simp only [eqToHom_refl, Category.comp_id]; exact M.left_unit X
  right_unit X _ := by exact M.right_unit X

def toMonad (M : GradedMonad C) : Monad C :=
  { M 0 with
    η := M.η
    μ := M.μ 0 0
    assoc X := by
      have h := M.assoc X 0 0 0
      simp only [Functor.comp_obj, Nat.add_zero, eqToHom_refl, Category.comp_id] at h
      exact h
    left_unit X := by
      have h := M.left_unit X 0
      simp only [Functor.id_obj, Nat.add_zero, eqToHom_refl, Category.comp_id] at h
      exact h
    right_unit X := M.right_unit X 0 }

lemma toMonad_ofMonad {M : Monad C} : (ofMonad M).toMonad = M := rfl

private def compFunctor (F : C ⥤ C) : ℕ → C ⥤ C
  | 0 => 𝟭 C
  | n + 1 => compFunctor F n ⋙ F

private lemma compFunctor_comp (F : C ⥤ C) (n k : ℕ) :
    compFunctor F k ⋙ compFunctor F n = compFunctor F (n + k) := by
  induction' n <;> all_goals expose_names
  · simp [compFunctor, Functor.comp_id]
  · rw [show n + 1 + k = n + k + 1 by ac_nf, compFunctor, compFunctor, ← h]; rfl

def ofComp (F : C ⥤ C) : GradedMonad C where
  toFunctor := compFunctor F
  η := NatTrans.id (𝟭 C)
  μ n k := NatTrans.id (compFunctor F k ⋙ compFunctor F n) ≫ eqToHom (compFunctor_comp _ n k)
  assoc X n k m := by simp [eqToHom_map]
  left_unit X n := by simp [compFunctor]
  right_unit X n := by simp [compFunctor]

end GradedMonad

end Instances

end CategoryTheory