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.Strong
open import Categories.Monad.Relative renaming (Monad to RMonad)
open import Categories.Monad.Construction.Kleisli
open import Data.Product using (_,_)
open import Categories.Functor.Core

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


-- K is the Initial Strong Pre-Elgot Monad

module Monad.Instance.K.StrongPreElgot {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 Algebra.Elgot.Free cocartesian
open import Algebra.Elgot.Stable distributive
open import Monad.PreElgot cocartesian
open import Monad.StrongPreElgot distributive
open import Monad.Instance.K distributive
open import Monad.Instance.K.Commutative distributive MK
open import Monad.Instance.K.Strong distributive MK
open import Monad.Instance.K.PreElgot distributive MK using (K-isPreElgot; K-isInitialPreElgot)
open import Category.Construction.PreElgotMonads cocartesian
open import Category.Construction.StrongPreElgotMonads distributive
open import Category.Construction.ElgotAlgebras cocartesian

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

K-isStrongPreElgot : IsStrongPreElgot KStrong
K-isStrongPreElgot = record 
  { preElgot = K-isPreElgot 
  ; strengthen-preserves = τ-comm 
  }

K-strongPreElgot : StrongPreElgotMonad
K-strongPreElgot = record 
  { SM = KStrong 
  ; isStrongPreElgot = K-isStrongPreElgot 
  }

isInitialStrongPreElgot : IsInitial StrongPreElgotMonads K-strongPreElgot
isInitialStrongPreElgot = record { ! = !′ ; !-unique = !-unique′ }
  where
    !′ : ∀ {A : StrongPreElgotMonad} → StrongPreElgotMonad-Morphism K-strongPreElgot A
    !′ {A} = record
      { α = PreElgotMonad-Morphism.α (IsInitial.! K-isInitialPreElgot {A-preElgot})
      ; α-η = α-η
      ; α-μ = α-μ
      ; α-strength = α-strength 
      ; α-preserves = λ {X} {B} f → Elgot-Algebra-Morphism.preserves (((freealgebras B) FreeObject.*) {A = record { A = T.F.F₀ B ; algebra = StrongPreElgotMonad.elgotalgebras A }} (T.η.η B)) 
      }
      where
        A-preElgot : PreElgotMonad
        A-preElgot = record { T = M ; isPreElgot = preElgot }
          where 
          open StrongPreElgotMonad A using (SM; isStrongPreElgot)
          open StrongMonad SM using (M)
          open IsStrongPreElgot isStrongPreElgot using (preElgot)
        open StrongPreElgotMonad A using (SM)
        module SM = StrongMonad SM
        open SM using (strengthen) renaming (M to T)
        open RMonad (Monad⇒Kleisli C T) using (extend)
        open monadK using () renaming (η to ηK; μ to μK)
        open strongK using () renaming (strengthen to strengthenK)
        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 = StrongPreElgotMonad.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 → StrongPreElgotMonad.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))

        -- The 'pre-Elgot monad morphism'-part is inherited:
        commute : ∀ {X Y : Obj} (f : X ⇒ Y) → ‼ Y ∘ K.₁ f ≈ T.F.₁ f ∘ ‼ X
        commute = NaturalTransformation.commute (PreElgotMonad-Morphism.α (IsInitial.! K-isInitialPreElgot {A-preElgot}))
        α-η : ∀ {X : Obj} → ‼ X ∘ ηK.η X ≈ T.η.η X
        α-η {X} = PreElgotMonad-Morphism.α-η (IsInitial.! K-isInitialPreElgot {A-preElgot})
        α-μ : ∀ {X : Obj} → ‼ X ∘ μK.η X ≈ T.μ.η X ∘ T.F.₁ (‼ X) ∘ ‼ (K.₀ X)
        α-μ {X} = PreElgotMonad-Morphism.α-μ (IsInitial.! K-isInitialPreElgot {A-preElgot})

        -- We are left to check strength, which follows by stability:

        α-strength : ∀ {X Y : Obj} → ‼ (X × Y) ∘ strengthenK.η (X , Y) ≈ strengthen.η (X , Y) ∘ (id ⁂ ‼ Y)
        α-strength {X} {Y} = by-stability (T-Alg (X × Y)) (T.η.η (X × Y)) (sym law₁) (sym law₂) pres₁ pres₂
          where
          law₁ : (‼ (X × Y) ∘ strengthenK.η (X , Y)) ∘ (id ⁂ ηK.η Y) ≈ T.η.η (X × Y)
          law₁ = begin 
            (‼ (X × Y) ∘ strengthenK.η (X , Y)) ∘ (id ⁂ ηK.η Y) ≈⟨ pullʳ (τ-η (X , Y)) ⟩ 
            ‼ (X × Y) ∘ ηK.η (X × Y)                             ≈⟨ α-η ⟩ 
            T.η.η (X × Y)                                         ∎
          law₂ : (strengthen.η (X , Y) ∘ (id ⁂ ‼ Y)) ∘ (id ⁂ ηK.η Y) ≈ T.η.η (X × Y)
          law₂ = begin 
            (strengthen.η (X , Y) ∘ (id ⁂ ‼ Y)) ∘ (id ⁂ ηK.η Y) ≈⟨ pullʳ (⁂∘⁂ ○ ⁂-cong₂ identity² α-η) ⟩ 
            strengthen.η (X , Y) ∘ (id ⁂ T.η.η Y)                 ≈⟨ SM.η-comm ⟩ 
            T.η.η (X × Y)                                          ∎
          pres₁ : ∀ {Z : Obj} (h : Z ⇒ K.₀ Y + Z) → (‼ (X × Y) ∘ strengthenK.η (X , Y)) ∘ (id ⁂ h #K) ≈ ((‼ (X × Y) ∘ strengthenK.η (X , Y) +₁ id) ∘ distributeˡ⁻¹ ∘ (id ⁂ h)) #T
          pres₁ {Z} h = begin 
            (‼ (X × Y) ∘ strengthenK.η (X , Y)) ∘ (id ⁂ h #K)                                   ≈⟨ pullʳ (τ-comm h) ⟩ 
            ‼ (X × Y) ∘ ((τ (X , Y) +₁ id) ∘ distributeˡ⁻¹ ∘ (id ⁂ h)) #K                      ≈⟨ ‼-preserves ((τ (X , Y) +₁ id) ∘ distributeˡ⁻¹ ∘ (id ⁂ h)) ⟩ 
            ((‼ (X × Y) +₁ id) ∘ (strengthenK.η (X , Y) +₁ id) ∘ distributeˡ⁻¹ ∘ (id ⁂ h)) #T ≈⟨ #-resp-≈ (StrongPreElgotMonad.elgotalgebras A) (pullˡ (+₁∘+₁ ○ +₁-cong₂ refl identity²)) ⟩ 
            ((‼ (X × Y) ∘ strengthenK.η (X , Y) +₁ id) ∘ distributeˡ⁻¹ ∘ (id ⁂ h)) #T          ∎
          pres₂ : ∀ {Z : Obj} (h : Z ⇒ K.₀ Y + Z) → (strengthen.η (X , Y) ∘ (id ⁂ ‼ Y)) ∘ (id ⁂ h #K) ≈ ((strengthen.η (X , Y) ∘ (id ⁂ ‼ Y) +₁ id) ∘ distributeˡ⁻¹ ∘ (id ⁂ h)) #T
          pres₂ {Z} h = begin 
            (strengthen.η (X , Y) ∘ (id ⁂ ‼ Y)) ∘ (id ⁂ h #K)                                       ≈⟨ pullʳ (⁂∘⁂ ○ ⁂-cong₂ identity² (‼-preserves h)) ⟩ 
            strengthen.η (X , Y) ∘ (id ⁂ ((‼ Y +₁ id) ∘ h) #T)                                      ≈⟨ StrongPreElgotMonad.strengthen-preserves A ((‼ Y +₁ id) ∘ h) ⟩ 
            ((strengthen.η (X , Y) +₁ id) ∘ distributeˡ⁻¹ ∘ (id ⁂ (‼ Y +₁ id) ∘ h)) #T             ≈⟨ sym (#-resp-≈ (StrongPreElgotMonad.elgotalgebras A) (refl⟩∘⟨ (pullʳ (⁂∘⁂ ○ ⁂-cong₂ identity² refl)))) ⟩ 
            (((strengthen.η (X , Y) +₁ id) ∘ (distributeˡ⁻¹ ∘ (id ⁂ (‼ Y +₁ id))) ∘ (id ⁂ h)) #T) ≈⟨ sym (#-resp-≈ (StrongPreElgotMonad.elgotalgebras A) (refl⟩∘⟨ (pullˡ ((+₁-cong₂ refl (sym (⟨⟩-unique id-comm id-comm))) ⟩∘⟨refl ○ distributeˡ⁻¹-natural id (‼ Y) id)))) ⟩ 
            ((strengthen.η (X , Y) +₁ id) ∘ ((id ⁂ ‼ Y) +₁ id) ∘ distributeˡ⁻¹ ∘ (id ⁂ h)) #T     ≈⟨ #-resp-≈ (StrongPreElgotMonad.elgotalgebras A) (pullˡ (+₁∘+₁ ○ +₁-cong₂ refl identity²)) ⟩ 
            ((strengthen.η (X , Y) ∘ (id ⁂ ‼ Y) +₁ id) ∘ distributeˡ⁻¹ ∘ (id ⁂ h)) #T              ∎

    -- ‼ is unique (somewhat different type profile than in PreElgot, so we reiterate the (short) proof)
    !-unique′ : ∀ {A : StrongPreElgotMonad} (f : StrongPreElgotMonad-Morphism K-strongPreElgot A) → StrongPreElgotMonad-Morphism.α (!′ {A = A}) ≃ StrongPreElgotMonad-Morphism.α f
    !-unique′ {A} f {X} = sym (FreeObject.*-uniq
                                (freealgebras X) 
                                {A = record { A = T.F.F₀ X ; algebra = StrongPreElgotMonad.elgotalgebras A }} 
                                (T.η.η X) 
                                (record { h = α.η X ; preserves = α-preserves _ }) 
                                α-η)
      where
        open StrongPreElgotMonad-Morphism f using (α; α-η; α-preserves)
        open StrongPreElgotMonad A using (SM)
        open StrongMonad SM using () renaming (M to T)