open import Level
open import Categories.Category.Core
open import Categories.FreeObjects.Free using (FreeObject; FO⇒Functor; FO⇒LAdj)
open import Categories.Functor.Core using (Functor)
open import Categories.Adjoint using (_⊣_)
open import Categories.Adjoint.Properties using (adjoint⇒monad)
open import Categories.Monad using (Monad)
open import Categories.Monad.Relative using () renaming (Monad to RMonad)
open import Categories.Monad.Construction.Kleisli
open import Categories.Category.Distributive using (Distributive)

import Categories.Morphism as M
import Categories.Morphism.Reasoning as MR

module Monad.Instance.K {o ℓ e} {C : Category o ℓ e} (distributive : Distributive C) where
  open import Category.Distributive.Helper distributive renaming (η to η-prod)
  open Category C
  open import Category.Construction.ElgotAlgebras cocartesian
  open import Algebra.Elgot cocartesian using (Elgot-Algebra)
  open import Algebra.Elgot.Free cocartesian using (FreeElgotAlgebra; elgotForgetfulF)
  open import Algebra.Elgot.Stable distributive using (IsStableFreeElgotAlgebra; IsStableFreeElgotAlgebraˡ; isStable⇒isStableˡ)


  open Equiv
  open MR C
  open M C
  open HomReasoning

  record MonadK : Set (suc o ⊔ suc ℓ ⊔ suc e) where
    field
      freealgebras : ∀ X → FreeElgotAlgebra X
      stable : ∀ X → IsStableFreeElgotAlgebra (freealgebras X)

    -- helper for accessing elgot algebras
    algebras : ∀ (X : Obj) → Elgot-Algebra
    algebras X = FreeObject.FX (freealgebras X)

    freeF : Functor C Elgot-Algebras
    freeF = FO⇒Functor elgotForgetfulF freealgebras
    
    adjoint : freeF ⊣ elgotForgetfulF
    adjoint = FO⇒LAdj elgotForgetfulF freealgebras

    monadK : Monad C
    monadK = adjoint⇒monad adjoint
    module monadK = Monad monadK


    kleisliK : KleisliTriple C
    kleisliK = Monad⇒Kleisli C monadK
    module kleisliK = RMonad kleisliK

    module K = Functor monadK.F

    open Elgot-Algebra using (#-resp-≈; #-Fixpoint; #-Compositionality; #-Uniformity; #-Folding; #-Diamond; #-Stutter) renaming (A to ⟦_⟧) public
    stableˡ = λ X → isStable⇒isStableˡ (freealgebras X) (stable X)
    η = λ Z → FreeObject.η (freealgebras Z)
    _♯ = λ {A X Y} f → IsStableFreeElgotAlgebra.[_,_]♯ {Y = X} (stable X) {X = A} (algebras Y) f
    _♯ˡ = λ {A X Y} f → IsStableFreeElgotAlgebraˡ.[_,_]♯ˡ {Y = X} (stableˡ X) {X = A} (algebras Y) f
    _# = λ {A} {X} f → Elgot-Algebra._# (algebras A) {X = X} f

    open kleisliK using (extend)
    open monadK using (μ)

    K-preserve : ∀ {X Y Z} (f : X ⇒ Y) (h : Z ⇒ K.₀ X + Z) → K.₁ f ∘ h # ≈ ((K.₁ f +₁ id) ∘ h) #
    K-preserve {X} {Y} {Z} f h = Elgot-Algebra-Morphism.preserves (((freealgebras _) FreeObject.*) (η _ ∘ f))
    
    extend-preserve : ∀ {X Y Z} (f : X ⇒ K.₀ Y) (h : Z ⇒ K.₀ X + Z) → extend f ∘ h # ≈ ((extend f +₁ id) ∘ h) #
    extend-preserve {X} {Y} {Z} f h = begin
      (μ.η _ ∘ K.₁ f) ∘ h #                 ≈⟨ pullʳ (Elgot-Algebra-Morphism.preserves (((freealgebras _) FreeObject.*) (η _ ∘ f))) ⟩
      μ.η _ ∘ ((K.₁ f +₁ id) ∘ h) #         ≈⟨ Elgot-Algebra-Morphism.preserves (((freealgebras _) FreeObject.*) id) ⟩
      ((μ.η _ +₁ id) ∘ (K.₁ f +₁ id) ∘ h) # ≈⟨ #-resp-≈ (algebras _) (pullˡ (+₁∘+₁ ○ +₁-cong₂ refl identity²)) ⟩
      ((extend f +₁ id) ∘ h) #              ∎

    μ-preserve : ∀ {X Y} (h : Y ⇒ K.₀ (K.₀ X) + Y) → μ.η X ∘ h # ≈ ((μ.η X +₁ id) ∘ h) #
    μ-preserve {X} {Y} h = Elgot-Algebra-Morphism.preserves (((freealgebras _) FreeObject.*) id)

    by-stability : ∀ {X Y} (A : Elgot-Algebra) {f g : X × ⟦ algebras Y ⟧ ⇒ ⟦ A ⟧} (i : X × Y ⇒ ⟦ A ⟧)
                    → i ≈ f ∘ (id ×₁ η Y) 
                    → i ≈ g ∘ (id ×₁ η Y) 
                    → (∀ {Z} (h : Z ⇒ ⟦ algebras Y ⟧ + Z) → f ∘ (id ×₁ (Elgot-Algebra._# (algebras Y) h)) ≈ Elgot-Algebra._# A ((f +₁ id) ∘ distributeˡ⁻¹ ∘ (id ×₁ h)))
                    → (∀ {Z} (h : Z ⇒ ⟦ algebras Y ⟧ + Z) → g ∘ (id ×₁ (Elgot-Algebra._# (algebras Y) h)) ≈ Elgot-Algebra._# A ((g +₁ id) ∘ distributeˡ⁻¹ ∘ (id ×₁ h)))
                    → f ≈ g
    by-stability {X} {Y} A {f} {g} i f-law g-law f-pres g-pres = begin 
      f                   ≈⟨ ♯-unique i f f-law f-pres ⟩ 
      [ A , i ]♯ ≈⟨ sym (♯-unique i g g-law g-pres) ⟩ 
      g                   ∎
      where
      open IsStableFreeElgotAlgebra (stable Y) using ([_,_]♯; ♯-unique)

    by-stabilityˡ : ∀ {X Y} (A : Elgot-Algebra) {f g : ⟦ algebras Y ⟧ × X ⇒ ⟦ A ⟧} (i : Y × X ⇒ ⟦ A ⟧) 
                    → i ≈ f ∘ (η Y ×₁ id) 
                    → i ≈ g ∘ (η Y ×₁ id) 
                    → (∀ {Z} (h : Z ⇒ ⟦ algebras Y ⟧ + Z) → f ∘ ((Elgot-Algebra._# (algebras Y) h) ×₁ id) ≈ Elgot-Algebra._# A ((f +₁ id) ∘ distributeʳ⁻¹ ∘ (h ×₁ id)))
                    → (∀ {Z} (h : Z ⇒ ⟦ algebras Y ⟧ + Z) → g ∘ ((Elgot-Algebra._# (algebras Y) h) ×₁ id) ≈ Elgot-Algebra._# A ((g +₁ id) ∘ distributeʳ⁻¹ ∘ (h ×₁ id)))
                    → f ≈ g
    by-stabilityˡ {X} {Y} A {f} {g} i f-law g-law f-pres g-pres = begin 
      f           ≈⟨ ♯ˡ-unique i f f-law f-pres ⟩ 
      [ A , i ]♯ˡ ≈⟨ sym (♯ˡ-unique i g g-law g-pres) ⟩ 
      g           ∎
      where
      open IsStableFreeElgotAlgebraˡ (stableˡ Y) using ([_,_]♯ˡ; ♯ˡ-unique)