Monad.lean

import Mathlib.tactic

open CategoryTheory

@[ext]
structure Mnd (X : Type) [Category X] extends T : X ⥤ X where
  η : Functor.id X ⟶ T
  μ : T ⋙ T ⟶ T
  assoc (x : X) : T.map (μ.app x) ≫ μ.app x = μ.app (T.obj x) ≫ μ.app x
  left_unit (x : X) : η.app (T.obj x) ≫ μ.app x = 𝟙 (T.obj x)
  right_unit (x : X) : T.map (η.app x) ≫ μ.app x = 𝟙 (T.obj x)


def Adjunction.Mnd [Category X] [Category A] {F : X ⥤ A} {G : A ⥤ X} (φ : F ⊣ G) : Mnd X := by
  use F ⋙ G, φ.unit, whiskerLeft F (whiskerRight φ.counit G)
  · intro x; simp
    rw [<- G.map_comp, <- G.map_comp]
    have E := φ.counit.naturality (φ.counit.app (F.obj x))
    rw [Functor.comp_map] at E
    conv at E => arg 1; arg 2; simp
    rw [E]
    simp
  · intro x
    conv => arg 1; arg 1; simp
    conv => arg 1; arg 2; simp
    conv => arg 2; simp
    rw [φ.right_triangle_components]
  · intro x; simp
    rw [<- G.map_comp, φ.left_triangle_components]; simp

structure Mnd.Alg [Category X] (M : Mnd X) where
  A : X
  a : M.T.obj A ⟶ A
  assoc : M.μ.app A ≫ a = M.T.map a ≫ a
  unit : M.η.app A ≫ a = 𝟙 A

@[ext]
structure Mnd.AlgHom [Category X] {M : Mnd X} (x x' : Mnd.Alg M) where
  hom : x.A ⟶ x'.A
  w : x.a ≫ hom = M.T.map hom ≫ x'.a

instance Mnd.Alg.category [Category X] (M : Mnd X) : Category (Mnd.Alg M) where
  Hom := Mnd.AlgHom
  id x := {
    hom := 𝟙 x.A
    w := by simp
  }
  comp {x y z} f g := by
    rcases f with ⟨f, Hf⟩
    rcases g with ⟨g, Hg⟩
    refine {
      hom := f ≫ g
      w := by simp; rw [<- Hg, <- Category.assoc, Hf]; simp
    }
  id_comp f := by simp
  comp_id f := by simp
  assoc f g h := by simp

def Mnd.free [Category X] (T : Mnd X) :
  X ⥤ Mnd.Alg T
:= {
  obj x := {
    A := T.obj x
    a := T.μ.app x
    assoc := by rw [T.assoc x]
    unit := by rw [T.left_unit x]
  }
  map {x x'} f := {
    hom := T.map f
    w := by
      simp
      have E := T.μ.naturality (f)
      rw [<- E]; simp
  }
  map_id x := by
    apply Mnd.AlgHom.ext;
    simp [CategoryStruct.id]
  map_comp {x y z} f g := by
    apply Mnd.AlgHom.ext
    simp [CategoryStruct.comp]
}

def Mnd.forget [Category X] (T : Mnd X) :
  Mnd.Alg T ⥤ X
:= {
  obj x := x.A
  map {x x'} f := f.hom
  map_id x := by simp [CategoryStruct.id]
  map_comp {x y z} f g := by
    rcases f with ⟨f, Hf⟩
    rcases g with ⟨g, Hg⟩
    simp [CategoryStruct.comp]
}

lemma Mnd.free_forget_eq [Category X] (T : Mnd X) :
  free T ⋙ forget T = T.T
:= by
  apply CategoryTheory.Functor.ext ?_ ?_
  · intro x; simp [free, forget]
  · intro x y f; simp [free, forget]

def Mnd.Adj [Category X] (T : Mnd X) :
  free T ⊣ forget T
:= {
  unit := T.η
  counit := {
    app x := {
      hom := x.a
      w := by
        simp [free, forget]
        rw [x.assoc]
    }
    naturality {x y} f := by
      rcases f with ⟨f, Hf⟩
      simp [free, forget]
      apply Mnd.AlgHom.ext
      simp [CategoryStruct.comp]
      rw [Hf]
  }
  left_triangle_components x := by
    simp [free, forget]
    simp [CategoryStruct.id, CategoryStruct.comp]
    apply Mnd.AlgHom.ext; simp
    have E := T.right_unit x
    rw [E]
  right_triangle_components x := by
    simp [free, forget]
    rw [x.unit]
}

def Mnd.Adj_eq_Mnd [Category X] (M : Mnd X) :
  Adjunction.Mnd M.Adj = M
:= by
  ext x
  · simp [Adjunction.Mnd, free, forget]
  · simp [Adjunction.Mnd, free, forget, Functor.comp]
  · simp [Adjunction.Mnd, free, forget, Mnd.Adj]
  · simp [Adjunction.Mnd, forget, free, Mnd.Adj]
    ext x; simp