open import Level
open import Categories.Category.Core
open import Categories.Monad.Relative using () renaming (Monad to RMonad)
open import Categories.Category.Cocartesian
open import Monad.Instance.Delay

import Categories.Morphism.Reasoning as MR

-- guarded morphisms of the delay monad can be uniformly iterated

module Monad.Instance.Delay.Guarded {o ℓ e} {C : Category o ℓ e} (cocartesian : Cocartesian C) (D : DelayM cocartesian) where
  open Category C
  open Cocartesian cocartesian
  open MR C
  open Equiv
  open HomReasoning
  open DelayM D
  open D-Kleisli
  open Coit

  record IsGuarded {X} {Y} (g : X ⇒ D₀ (Y + X)) : Set (ℓ ⊔ e) where
    field
      guard : X ⇒ Y + D₀ (Y + X)
      guard-eq : (i₁ +₁ id) ∘ guard ≈ out ∘ g

  -- guarded iteration
  _#[_] : ∀ {X Y} (g : X ⇒ D₀ (Y + X)) → IsGuarded g → X ⇒ D₀ Y
  _#[_] {X} {Y} g isGuarded = coit f ∘ g
    where
      open IsGuarded isGuarded
      f : D₀ (Y + X) ⇒ Y + D₀ (Y + X)
      f = [ [ i₁ , guard ] , i₂ ] ∘ out

  guarded-fixpoint : ∀ {X Y} (g : X ⇒ D₀ (Y + X)) (isGuarded : IsGuarded g) → g #[ isGuarded ] ≈ extend ([ now , g #[ isGuarded ] ]) ∘ g
  guarded-fixpoint {X} {Y} g isGuarded = (sym (extend-unique [ now , g #[ isGuarded ] ] (coit ([ [ i₁ , guard ] , i₂ ] ∘ out)) (begin 
    out ∘ coit ([ [ i₁ , guard ] , i₂ ] ∘ out)                                                                                                                                 ≈⟨ coit-commutes ([ [ i₁ , guard ] , i₂ ] ∘ out) ⟩ 
    (id +₁ coit ([ [ i₁ , guard ] , i₂ ] ∘ out)) ∘ [ [ i₁ , guard ] , i₂ ] ∘ out                                                                                               ≈⟨ pullˡ (∘[] ○ []-cong₂ ∘[] refl) ⟩ 
    [ [ (id +₁ coit ([ [ i₁ , guard ] , i₂ ] ∘ out)) ∘ i₁ , (id +₁ coit ([ [ i₁ , guard ] , i₂ ] ∘ out)) ∘ guard ] , (id +₁ coit ([ [ i₁ , guard ] , i₂ ] ∘ out)) ∘ i₂ ] ∘ out ≈⟨ ([]-cong₂ ([]-cong₂ (inject₁ ○ identityʳ) refl) inject₂) ⟩∘⟨refl ⟩ 
    [ [ i₁ , (id +₁ coit ([ [ i₁ , guard ] , i₂ ] ∘ out)) ∘ guard ] , i₂ ∘ coit ([ [ i₁ , guard ] , i₂ ] ∘ out) ] ∘ out                                                        ≈˘⟨ ([]-cong₂ ([]-cong₂ refl (refl⟩∘⟨ (cancelˡ ([]∘+₁ ○ []-cong₂ inject₁ identityʳ ○ +-η)))) refl) ⟩∘⟨refl ⟩ 
    [ [ i₁ , (id +₁ coit ([ [ i₁ , guard ] , i₂ ] ∘ out)) ∘ [ [ i₁ , guard ] , i₂ ] ∘ (i₁ +₁ id) ∘ guard ] , i₂ ∘ coit ([ [ i₁ , guard ] , i₂ ] ∘ out) ] ∘ out                 ≈˘⟨ ([]-cong₂ ([]-cong₂ refl (pullʳ (pullʳ (sym guard-eq)))) refl) ⟩∘⟨refl ⟩ 
    [ [ i₁ , ((id +₁ coit ([ [ i₁ , guard ] , i₂ ] ∘ out)) ∘ [ [ i₁ , guard ] , i₂ ] ∘ out) ∘ g ] , i₂ ∘ coit ([ [ i₁ , guard ] , i₂ ] ∘ out) ] ∘ out                          ≈˘⟨ ([]-cong₂ ([]-cong₂ unitlaw (pullˡ (coit-commutes ([ [ i₁ , guard ] , i₂ ] ∘ out)))) refl) ⟩∘⟨refl ⟩ 
    [ [ out ∘ now , out ∘ g #[ isGuarded ] ] , i₂ ∘ coit ([ [ i₁ , guard ] , i₂ ] ∘ out) ] ∘ out                                                                               ≈˘⟨ ([]-cong₂ ∘[] refl) ⟩∘⟨refl ⟩ 
    [ out ∘ [ now , g #[ isGuarded ] ] , i₂ ∘ coit ([ [ i₁ , guard ] , i₂ ] ∘ out) ] ∘ out                                                                                     ∎))) ⟩∘⟨refl
    where
    open IsGuarded isGuarded

  guarded-unique : ∀ {X Y} (g : X ⇒ D₀ (Y + X)) (f i : X ⇒ D₀ Y) → IsGuarded g → f ≈ extend ([ now , f ]) ∘ g → i ≈ extend ([ now , i ]) ∘ g → f ≈ i
  guarded-unique {X} {Y} g f i isGuarded eqf eqi = begin 
    f                                        ≈⟨ eqf ⟩ 
    extend ([ now , f ]) ∘ g                 ≈⟨ sym (coit-unique ([ [ i₁ , guard ] , i₂ ] ∘ out) (extend ([ now , f ])) (coit-unique-helper f eqf)) ⟩∘⟨refl ⟩ 
    coit ([ [ i₁ , guard ] , i₂ ] ∘ out) ∘ g ≈⟨ coit-unique ([ [ i₁ , guard ] , i₂ ] ∘ out) (extend ([ now , i ])) (coit-unique-helper i eqi) ⟩∘⟨refl ⟩ 
    extend ([ now , i ]) ∘ g                 ≈⟨ sym eqi ⟩ 
    i                                        ∎
    where
    open IsGuarded isGuarded
    coit-unique-helper : ∀ (f : X ⇒ D₀ Y) → f ≈ extend ([ now , f ]) ∘ g → out ∘ extend ([ now , f ]) ≈ (id +₁ extend ([ now , f ])) ∘ [ [ i₁ , guard ] , i₂ ] ∘ out
    coit-unique-helper f eqf = begin 
      out ∘ extend ([ now , f ])                                                                                         ≈⟨ extend-commutes ([ now , f ]) ⟩ 
      [ out ∘ [ now , f ] , i₂ ∘ extend ([ now , f ]) ] ∘ out                                                            ≈⟨ ([]-cong₂ ∘[] refl) ⟩∘⟨refl ⟩ 
      [ [ out ∘ now , out ∘ f ] , i₂ ∘ extend ([ now , f ]) ] ∘ out                                                      ≈⟨ ([]-cong₂ ([]-cong₂ unitlaw out∘f) refl) ⟩∘⟨refl ⟩
      [ [ i₁ , (id +₁ extend ([ now , f ])) ∘ guard ] , i₂ ∘ extend ([ now , f ]) ] ∘ out                                ≈˘⟨ ([]-cong₂ ([]-cong₂ identityʳ refl) refl) ⟩∘⟨refl ⟩
      [ [ i₁ ∘ id , (id +₁ extend ([ now , f ])) ∘ guard ] , i₂ ∘ extend ([ now , f ]) ] ∘ out                           ≈˘⟨ ([]-cong₂ ([]-cong₂ +₁∘i₁ refl) refl) ⟩∘⟨refl ⟩ 
      [ [ (id +₁ extend ([ now , f ])) ∘ i₁ , (id +₁ extend ([ now , f ])) ∘ guard ] , i₂ ∘ extend ([ now , f ]) ] ∘ out ≈˘⟨ ([]-cong₂ ∘[] +₁∘i₂) ⟩∘⟨refl ⟩ 
      [ (id +₁ extend ([ now , f ])) ∘ [ i₁ , guard ] , (id +₁ extend ([ now , f ])) ∘ i₂ ] ∘ out                        ≈˘⟨ pullˡ ∘[] ⟩ 
      (id +₁ extend ([ now , f ])) ∘ [ [ i₁ , guard ] , i₂ ] ∘ out                                                       ∎
      where
      out∘f : out ∘ f ≈ (id +₁ extend ([ now , f ])) ∘ guard
      out∘f = begin 
        out ∘ f                                                                 ≈⟨ refl⟩∘⟨ eqf ⟩ 
        out ∘ extend ([ now , f ]) ∘ g                                          ≈⟨ pullˡ (extend-commutes ([ now , f ])) ⟩
        ([ out ∘ [ now , f ] , i₂ ∘ extend ([ now , f ]) ] ∘ out) ∘ g           ≈⟨ pullʳ (sym guard-eq) ⟩
        [ out ∘ [ now , f ] , i₂ ∘ extend ([ now , f ]) ] ∘ (i₁ +₁ id) ∘ guard  ≈⟨ pullˡ []∘+₁ ⟩
        [ (out ∘ [ now , f ]) ∘ i₁ , (i₂ ∘ extend ([ now , f ])) ∘ id ] ∘ guard ≈⟨ ([]-cong₂ (pullʳ inject₁) identityʳ) ⟩∘⟨refl ⟩
        [ out ∘ now , i₂ ∘ extend ([ now , f ]) ] ∘ guard                       ≈⟨ ([]-cong₂ (unitlaw ○ sym identityʳ) refl) ⟩∘⟨refl ⟩
        (id +₁ extend ([ now , f ])) ∘ guard                                    ∎


  guarded-uniformity : ∀ {X Y Z} (g : X ⇒ D₀ (Y + X)) (g-guarded : IsGuarded g) (h : X ⇒ Z) (f : Z ⇒ D₀ (Y + Z)) (f-guarded : IsGuarded f) → f ∘ h ≈ extend (now ∘ (id +₁ h)) ∘ g → g #[ g-guarded ] ≈ f #[ f-guarded ] ∘ h
  guarded-uniformity {X} {Y} {Z} g g-guarded h f f-guarded eq = guarded-unique g (g #[ g-guarded ]) (f #[ f-guarded ] ∘ h) g-guarded (guarded-fixpoint g g-guarded) (begin 
    f #[ f-guarded ] ∘ h                                               ≈⟨ (guarded-fixpoint f f-guarded) ⟩∘⟨refl ⟩
    (extend ([ now , f #[ f-guarded ] ]) ∘ f) ∘ h                      ≈⟨ pullʳ eq ⟩
    extend ([ now , f #[ f-guarded ] ]) ∘ extend (now ∘ (id +₁ h)) ∘ g ≈⟨ pullˡ (sym DK.assoc) ⟩
    extend (extend ([ now , f #[ f-guarded ] ]) ∘ now ∘ (id +₁ h)) ∘ g ≈⟨ extend-≈ (pullˡ DK.identityʳ) ⟩∘⟨refl ⟩
    extend ([ now , f #[ f-guarded ] ] ∘ (id +₁ h)) ∘ g                ≈⟨ (extend-≈ ([]∘+₁ ○ []-cong₂ identityʳ refl)) ⟩∘⟨refl ⟩
    extend ([ now , f #[ f-guarded ] ∘ h ]) ∘ g                        ∎)