{-# OPTIONS --without-K --safe #-}

open import Categories.Bicategory

module Categories.Bicategory.Monad.Bimodule {o  e t} {𝒞 : Bicategory o  e t} where

open import Level
open import Categories.Bicategory.Monad using (Monad)
import Categories.Bicategory.Extras as Bicat
open Bicat 𝒞
open Shorthands
import Categories.Morphism.Reasoning as MR

private
  module MR' {A B : Obj} where
    open MR (hom A B) using (pullˡ; elimʳ; assoc²βγ) public

record Bimodule (M₁ M₂ : Monad 𝒞) : Set (o    e) where
  open Monad using (C; T; μ; η)
  field
    F       : C M₁ ⇒₁ C M₂
    actionˡ : F ∘₁ T M₁ ⇒₂ F
    actionʳ : T M₂ ∘₁ F ⇒₂ F

    assoc     : actionʳ ∘ᵥ (T M₂  actionˡ) ∘ᵥ α⇒  actionˡ ∘ᵥ (actionʳ  T M₁)
    sym-assoc : actionˡ ∘ᵥ (actionʳ  T M₁) ∘ᵥ α⇐  actionʳ ∘ᵥ (T M₂  actionˡ)

    assoc-actionˡ     : actionˡ ∘ᵥ (F  μ M₁) ∘ᵥ α⇒  actionˡ ∘ᵥ (actionˡ  T M₁)
    sym-assoc-actionˡ : actionˡ ∘ᵥ (actionˡ  T M₁) ∘ᵥ α⇐  actionˡ ∘ᵥ (F  μ M₁)
    assoc-actionʳ     : actionʳ ∘ᵥ (μ M₂  F) ∘ᵥ α⇐  actionʳ ∘ᵥ (T M₂  actionʳ)
    sym-assoc-actionʳ : actionʳ ∘ᵥ (T M₂  actionʳ) ∘ᵥ α⇒  actionʳ ∘ᵥ (μ M₂  F)

    identityˡ : actionˡ ∘ᵥ (F  η M₁) ∘ᵥ unitorʳ.to  id₂
    identityʳ : actionʳ ∘ᵥ (η M₂  F) ∘ᵥ unitorˡ.to  id₂

-- This helper definition lets us specify only one of each associativity laws
-- and have the symmetric one derived.
record BIMODHelper (M₁ M₂ : Monad 𝒞) : Set (o    e) where
  open Monad using (C; T; μ; η)
  field
    F       : C M₁ ⇒₁ C M₂
    actionˡ : F ∘₁ T M₁ ⇒₂ F
    actionʳ : T M₂ ∘₁ F ⇒₂ F

    assoc : actionʳ ∘ᵥ (T M₂  actionˡ) ∘ᵥ α⇒  actionˡ ∘ᵥ (actionʳ  T M₁)

    assoc-actionˡ : actionˡ ∘ᵥ (F  μ M₁) ∘ᵥ α⇒  actionˡ ∘ᵥ (actionˡ  T M₁)
    assoc-actionʳ : actionʳ ∘ᵥ (μ M₂  F) ∘ᵥ α⇐  actionʳ ∘ᵥ (T M₂  actionʳ)

    identityˡ : actionˡ ∘ᵥ (F  η M₁) ∘ᵥ unitorʳ.to  id₂
    identityʳ : actionʳ ∘ᵥ (η M₂  F) ∘ᵥ unitorˡ.to  id₂

bimodHelper :  {M₁ M₂ : Monad 𝒞}  BIMODHelper M₁ M₂  Bimodule M₁ M₂
bimodHelper {M₁} {M₂} B = record
  { F = F
  ; actionˡ = actionˡ
  ; actionʳ = actionʳ
  ; assoc = assoc
  ; sym-assoc = sym-assoc
  ; assoc-actionˡ = assoc-actionˡ
  ; sym-assoc-actionˡ = sym-assoc-actionˡ
  ; assoc-actionʳ = assoc-actionʳ
  ; sym-assoc-actionʳ = sym-assoc-actionʳ
  ; identityˡ = identityˡ
  ; identityʳ = identityʳ
  }
  where
    open Monad using (T; μ)
    open BIMODHelper B
    open hom.HomReasoning
    open MR'

    sym-assoc : actionˡ ∘ᵥ (actionʳ  T M₁) ∘ᵥ α⇐  actionʳ ∘ᵥ (T M₂  actionˡ)
    sym-assoc = begin
      actionˡ ∘ᵥ (actionʳ  T M₁) ∘ᵥ α⇐         ≈⟨ pullˡ ( assoc) 
      (actionʳ ∘ᵥ (T M₂  actionˡ) ∘ᵥ α⇒) ∘ᵥ α⇐ ≈⟨ assoc²βγ 
      (actionʳ ∘ᵥ (T M₂  actionˡ)) ∘ᵥ α⇒ ∘ᵥ α⇐ ≈⟨ elimʳ associator.isoʳ 
      actionʳ ∘ᵥ (T M₂  actionˡ)               

    sym-assoc-actionˡ : actionˡ ∘ᵥ (actionˡ  T M₁) ∘ᵥ α⇐  actionˡ ∘ᵥ (F  μ M₁)
    sym-assoc-actionˡ = begin
      actionˡ ∘ᵥ (actionˡ  T M₁) ∘ᵥ α⇐   ≈⟨ pullˡ ( assoc-actionˡ) 
      (actionˡ ∘ᵥ (F  μ M₁) ∘ᵥ α⇒) ∘ᵥ α⇐ ≈⟨ assoc²βγ 
      (actionˡ ∘ᵥ (F  μ M₁)) ∘ᵥ α⇒ ∘ᵥ α⇐ ≈⟨ elimʳ associator.isoʳ 
      actionˡ ∘ᵥ (F  μ M₁)               

    sym-assoc-actionʳ : actionʳ ∘ᵥ (T M₂  actionʳ) ∘ᵥ α⇒  actionʳ ∘ᵥ (μ M₂  F)
    sym-assoc-actionʳ = begin
      actionʳ ∘ᵥ (T M₂  actionʳ) ∘ᵥ α⇒   ≈⟨ pullˡ ( assoc-actionʳ) 
      (actionʳ ∘ᵥ (μ M₂  F) ∘ᵥ α⇐) ∘ᵥ α⇒ ≈⟨ assoc²βγ 
      (actionʳ ∘ᵥ (μ M₂  F)) ∘ᵥ α⇐ ∘ᵥ α⇒ ≈⟨ elimʳ associator.isoˡ 
      actionʳ ∘ᵥ (μ M₂  F)               

id-bimodule : (M : Monad 𝒞)  Bimodule M M
id-bimodule M = record
  { F = T
  ; actionˡ = μ
  ; actionʳ = μ
  ; assoc = assoc
  ; sym-assoc = sym-assoc
  ; assoc-actionˡ = assoc
  ; sym-assoc-actionˡ = sym-assoc
  ; assoc-actionʳ = sym-assoc
  ; sym-assoc-actionʳ = assoc
  ; identityˡ = identityˡ
  ; identityʳ = identityʳ
  }
  where
    open Monad M