open import Level
open import Categories.Category.Core
open import Categories.Category.Distributive
open import Categories.FreeObjects.Free
open import Categories.Object.Initial
open import Categories.NaturalTransformation hiding (id)
open import Categories.NaturalTransformation.Equivalence
open import Categories.Monad hiding (id)
open import Categories.Monad.Relative renaming (Monad to RMonad)
open import Categories.Monad.Construction.Kleisli
open import Monad.Helper

import Categories.Morphism as M
import Categories.Morphism.Reasoning as MR
import Monad.Instance.K as MIK


-- K is the Initial Pre-Elgot Monad

module Monad.Instance.K.PreElgot {o ℓ e} {C : Category o ℓ e} (distributive : Distributive C) (MK : MIK.MonadK distributive) where
open import Category.Distributive.Helper distributive
open MIK distributive
open MonadK MK
open import Algebra.Elgot cocartesian
open import Monad.PreElgot cocartesian
open import Monad.Instance.K distributive
open import Monad.Instance.K.Commutative distributive MK
open import Monad.Instance.K.Strong distributive MK
open import Category.Construction.PreElgotMonads cocartesian
open import Category.Construction.ElgotAlgebras cocartesian

open Category C
open Equiv
open HomReasoning
open MR C
open M C

K-isPreElgot : IsPreElgot monadK
K-isPreElgot = record
  { elgotalgebras = λ {X} → Elgot-Algebra.algebra (algebras X)
  ; extend-preserves = λ f h → sym (extend-preserve h f)
  }
  where open kleisliK using (extend)

K-preElgot : PreElgotMonad
K-preElgot = record { T = monadK ; isPreElgot = K-isPreElgot }

K-isInitialPreElgot : IsInitial PreElgotMonads K-preElgot
K-isInitialPreElgot = record
  { ! = !′
  ; !-unique = !-unique′
  }
  where
    !′ : ∀ {A : PreElgotMonad} → PreElgotMonad-Morphism K-preElgot A
    !′ {A} = record
      { α = ntHelper (record
        { η = ‼
        ; commute = commute
        })
      ; α-η = FreeObject.*-lift (freealgebras _) (T.η.η _)
      ; α-μ = α-μ
      ; preserves = λ {X} {B} f → Elgot-Algebra-Morphism.preserves (((freealgebras B) FreeObject.*) {A = record { A = T.F.F₀ B ; algebra = PreElgotMonad.elgotalgebras A }} (T.η.η B))
      }
      where
        open PreElgotMonad A using (T)
        open RMonad (Monad⇒Kleisli C T) using (extend)
        module T = Monad T
        open monadK using () renaming (η to ηK; μ to μK)
        open Elgot-Algebra-on using (#-resp-≈)

        -- Shorthand for the elgot algebra structure on TX
        T-Alg : ∀ (X : Obj) → Elgot-Algebra
        T-Alg X = record { A = T.F.₀ X ; algebra = PreElgotMonad.elgotalgebras A }

        -- Shorthand for the elgot algebra structure on KX
        K-Alg : ∀ (X : Obj) → Elgot-Algebra
        K-Alg X = record { A = K.₀ X ; algebra = Elgot-Algebra.algebra (algebras X) }

        -- The initiality morphism
        ‼ : ∀ (X : Obj) → K.₀ X ⇒ T.F.₀ X
        ‼ X = Elgot-Algebra-Morphism.h (_* {A = T-Alg X} (T.η.η X))
          where open FreeObject (freealgebras X)

        -- Shorthands for iteration operators of KX and TX
        _#K = λ {B} {C} f → Elgot-Algebra._# (FreeObject.FX (freealgebras C)) {B} f
        _#T = λ {B} {C} f → PreElgotMonad.elgotalgebras._# A {B} {C} f

        -- Shorthands for some preservation facts
        K₁-preserves : ∀ {X Y Z : Obj} (f : X ⇒ Y) (g : Z ⇒ K.₀ X + Z) → K.₁ f ∘ (g #K) ≈ ((K.₁ f +₁ id) ∘ g) #K
        K₁-preserves {X} {Y} {Z} f g = Elgot-Algebra-Morphism.preserves (((freealgebras X) FreeObject.*) {A = K-Alg Y} (ηK.η _ ∘ f))
        μK-preserves : ∀ {X Y : Obj} (g : Y ⇒ K.₀ (K.₀ X) + Y) → μK.η X ∘ g #K ≈ ((μK.η X +₁ id) ∘ g) #K
        μK-preserves {X} g = Elgot-Algebra-Morphism.preserves (((freealgebras (K.₀ X)) FreeObject.*) {A = K-Alg X} id)
        ‼-preserves : ∀ {X Y : Obj} (g : Y ⇒ K.₀ X + Y) → ‼ X ∘ g #K ≈ ((‼ X +₁ id) ∘ g) #T
        ‼-preserves {X} g = Elgot-Algebra-Morphism.preserves (((freealgebras X) FreeObject.*) {A = T-Alg X} (T.η.η X))

        -- Naturality of ‼
        commute : ∀ {X Y : Obj} (f : X ⇒ Y) → ‼ Y ∘ K.₁ f ≈ T.F.₁ f ∘ ‼ X
        commute {X} {Y} f = begin
          ‼ Y ∘ K.₁ f                                                    ≈⟨ FreeObject.*-uniq
                                                                                (freealgebras X)
                                                                                {A = T-Alg Y} 
                                                                                (T.F.₁ f ∘ T.η.η X) 
                                                                                (record { h = ‼ Y ∘ K.₁ f ; preserves = pres₁ }) 
                                                                                comm₁ ⟩
          Elgot-Algebra-Morphism.h (FreeObject._* (freealgebras X) {A = T-Alg Y} (T.F.₁ f ∘ T.η.η _)) ≈⟨ sym (FreeObject.*-uniq 
                                                                                  (freealgebras X)
                                                                                  {A = T-Alg Y} 
                                                                                  (T.F.₁ f ∘ T.η.η X) 
                                                                                  (record { h = T.F.₁ f ∘ ‼ X ; preserves = pres₂ }) 
                                                                                  (pullʳ (FreeObject.*-lift (freealgebras X) (T.η.η X)))) ⟩
          T.F.₁ f ∘ ‼ X                                                  ∎
          where
            pres₁ : ∀ {Z} {g : Z ⇒ K.₀ X + Z} → (‼ Y ∘ K.₁ f) ∘ g #K ≈ ((‼ Y ∘ K.₁ f +₁ id) ∘ g) #T
            pres₁ {Z} {g} = begin
              (‼ Y ∘ K.₁ f) ∘ (g #K)                   ≈⟨ pullʳ (K₁-preserves f g) ⟩
              ‼ Y ∘ (((K.₁ f +₁ id) ∘ g) #K)          ≈⟨ ‼-preserves ((K.₁ f +₁ id) ∘ g) ⟩
              (((‼ Y +₁ id) ∘ (K.₁ f +₁ id) ∘ g) #T) ≈⟨ #-resp-≈ (PreElgotMonad.elgotalgebras A) (pullˡ (+₁∘+₁ ○ +₁-cong₂ refl identity²)) ⟩
              ((‼ Y ∘ K.₁ f +₁ id) ∘ g) #T            ∎
            pres₂ : ∀ {Z} {g : Z ⇒ K.₀ X + Z} → (T.F.₁ f ∘ ‼ X) ∘ g #K ≈ ((T.F.₁ f ∘ ‼ X +₁ id) ∘ g) #T
            pres₂ {Z} {g} = begin
              (T.F.₁ f ∘ ‼ X) ∘ g #K                                  ≈⟨ pullʳ (‼-preserves g) ⟩
              T.F.₁ f ∘ ((‼ X +₁ id) ∘ g) #T                         ≈⟨ (sym (F₁⇒extend T f)) ⟩∘⟨refl ⟩
              extend (T.η.η Y ∘ f) ∘ ((‼ X +₁ id) ∘ g) #T            ≈⟨ sym (PreElgotMonad.extend-preserves A ((‼ X +₁ id) ∘ g) (T.η.η Y ∘ f)) ⟩
              (((extend (T.η.η Y ∘ f) +₁ id) ∘ (‼ X +₁ id) ∘ g) #T) ≈⟨ #-resp-≈ (PreElgotMonad.elgotalgebras A) (pullˡ (+₁∘+₁ ○ +₁-cong₂ ((F₁⇒extend T f) ⟩∘⟨refl) identity²)) ⟩
              ((T.F.₁ f ∘ ‼ X +₁ id) ∘ g) #T                         ∎
            comm₁ : (‼ Y ∘ K.₁ f) ∘ _ ≈ T.F.₁ f ∘ T.η.η X
            comm₁ = begin
              (‼ Y ∘ K.₁ f) ∘ _ ≈⟨ pullʳ (K₁η f) ⟩
              ‼ Y ∘ ηK.η _ ∘ f  ≈⟨ pullˡ (FreeObject.*-lift (freealgebras Y) (T.η.η Y)) ⟩
              T.η.η Y ∘ f        ≈⟨ NaturalTransformation.commute T.η f ⟩
              T.F.₁ f ∘ T.η.η X  ∎

        -- ‼ respects multiplication
        α-μ : ∀ {X : Obj} → ‼ X ∘ μK.η X ≈ T.μ.η X ∘ T.F.₁ (‼ X) ∘ ‼ (K.₀ X)
        α-μ {X} = begin
          ‼ X ∘ μK.η X                                                                      ≈⟨ FreeObject.*-uniq 
                            (freealgebras (K.₀ X)) 
                            {A = T-Alg X} 
                            (‼ X) 
                            (record { h = ‼ X ∘ μK.η X ; preserves = pres₁ }) 
                            (cancelʳ monadK.identityʳ) ⟩
          Elgot-Algebra-Morphism.h (((freealgebras (K.₀ X)) FreeObject.*) {A = T-Alg X} (‼ X)) ≈⟨ sym (FreeObject.*-uniq 
                                                                                                      (freealgebras (K.₀ X)) 
                                                                                                      {A = T-Alg X} 
                                                                                                      (‼ X) 
                                                                                                      (record { h = T.μ.η X ∘ T.F.₁ (‼ X) ∘ ‼ (K.₀ X) ; preserves = pres₂ }) 
                                                                                                      comm) ⟩
          T.μ.η X ∘ T.F.₁ (‼ X) ∘ ‼ (K.₀ X)                                                ∎
          where
            pres₁ : ∀ {Z} {g : Z ⇒ K.₀ (K.₀ X) + Z} → (‼ X ∘ μK.η X) ∘ g #K ≈ ((‼ X ∘ μK.η X +₁ id) ∘ g) #T
            pres₁ {Z} {g} = begin
              (‼ X ∘ μK.η X) ∘ (g #K)                   ≈⟨ pullʳ (μK-preserves g) ⟩
              ‼ X ∘ ((μK.η X +₁ id) ∘ g) #K            ≈⟨ ‼-preserves ((μK.η X +₁ id) ∘ g) ⟩
              (((‼ X +₁ id) ∘ (μK.η X +₁ id) ∘ g) #T) ≈⟨ #-resp-≈ (PreElgotMonad.elgotalgebras A) (pullˡ (+₁∘+₁ ○ +₁-cong₂ refl identity²)) ⟩
              (((‼ X ∘ μK.η X +₁ id) ∘ g) #T)          ∎
            pres₂ : ∀ {Z} {g : Z ⇒ K.₀ (K.₀ X) + Z} → (T.μ.η X ∘ T.F.₁ (‼ X) ∘ ‼ (K.₀ X)) ∘ g #K ≈ ((T.μ.η X ∘ T.F.₁ (‼ X) ∘ ‼ (K.₀ X) +₁ id) ∘ g) #T
            pres₂ {Z} {g} = begin
              (T.μ.η X ∘ T.F.₁ (‼ X) ∘ ‼ (K.₀ X)) ∘ (g #K)                                      ≈⟨ pullʳ (pullʳ (‼-preserves g)) ⟩
              T.μ.η X ∘ T.F.₁ (‼ X) ∘ (((‼ (K.₀ X) +₁ id) ∘ g) #T)                             ≈⟨ refl⟩∘⟨ ((sym (F₁⇒extend T (‼ X))) ⟩∘⟨refl ○ sym (PreElgotMonad.extend-preserves A ((‼ (K.₀ X) +₁ id) ∘ g) (T.η.η (T.F.F₀ X) ∘ ‼ X)) )⟩
              T.μ.η X ∘ ((extend (T.η.η _ ∘ ‼ _) +₁ id) ∘ ((‼ _ +₁ id)) ∘ g) #T               ≈⟨ (sym (elimʳ T.F.identity)) ⟩∘⟨refl ⟩
              extend id ∘ ((extend (T.η.η _ ∘ ‼ _) +₁ id) ∘ ((‼ _ +₁ id)) ∘ g) #T            ≈⟨ sym (PreElgotMonad.extend-preserves A ((extend (T.η.η (T.F.F₀ X) ∘ ‼ X) +₁ id) ∘ (‼ (K.₀ X) +₁ id) ∘ g) id) ⟩
              (((extend id +₁ id) ∘ (extend (T.η.η _ ∘ ‼ _) +₁ id) ∘ ((‼ _ +₁ id)) ∘ g) #T) ≈⟨ #-resp-≈ (PreElgotMonad.elgotalgebras A) (pullˡ (+₁∘+₁ ○ +₁-cong₂ ((elimʳ T.F.identity) ⟩∘⟨ (F₁⇒extend T (‼ X))) identity²)) ⟩
              (((T.μ.η X ∘ T.F.₁ (‼ X) +₁ id) ∘ (‼ _ +₁ id) ∘ g) #T)                          ≈⟨ #-resp-≈ (PreElgotMonad.elgotalgebras A) (pullˡ (+₁∘+₁ ○ +₁-cong₂ assoc identity²)) ⟩
              (((T.μ.η X ∘ T.F.₁ (‼ X) ∘ ‼ (K.₀ X) +₁ id) ∘ g) #T)                             ∎
            comm : (T.μ.η X ∘ T.F.₁ (‼ X) ∘ ‼ (K.₀ X)) ∘ ηK.η (K.₀ X) ≈ ‼ X
            comm = begin
              (T.μ.η X ∘ T.F.₁ (‼ X) ∘ ‼ (K.₀ X)) ∘ ηK.η (K.₀ X) ≈⟨ (refl⟩∘⟨ sym (commute (‼ X))) ⟩∘⟨refl ⟩
              (T.μ.η X ∘ ‼ _ ∘ K.₁ (‼ X)) ∘ ηK.η (K.₀ X)         ≈⟨ assoc ○ refl⟩∘⟨ (assoc ○  refl⟩∘⟨ sym (monadK.η.commute (‼ X))) ⟩
              T.μ.η X ∘ ‼ _ ∘ ηK.η (T.F.F₀ X) ∘ ‼ X              ≈⟨ refl⟩∘⟨ (pullˡ (FreeObject.*-lift (freealgebras _) (T.η.η _))) ⟩
              T.μ.η X ∘ T.η.η _ ∘ ‼ X                             ≈⟨ cancelˡ (Monad.identityʳ T) ⟩
              ‼ X                                                 ∎

    -- ‼ is unique
    !-unique′ : ∀ {A : PreElgotMonad} (f : PreElgotMonad-Morphism K-preElgot A) → PreElgotMonad-Morphism.α (!′ {A = A}) ≃ PreElgotMonad-Morphism.α f
    !-unique′ {A} f {X} = sym (FreeObject.*-uniq
                                (freealgebras X) 
                                {A = record { A = T.F.F₀ X ; algebra = PreElgotMonad.elgotalgebras A }} 
                                (T.η.η X) 
                                (record { h = α.η X ; preserves = preserves _ }) 
                                α-η)
      where
        open PreElgotMonad-Morphism f using (α; α-η; preserves)
        open PreElgotMonad A using (T)
        module T = Monad T