{-# OPTIONS --allow-unsolved-metas #-}

open import Categories.Category
open import Monad.Instance.Delay
open import Categories.Category.Distributive
open import Categories.Object.Terminal 
open import Categories.Object.NaturalNumbers.Parametrized
open import Categories.Category.Cartesian

open import Categories.Object.Exponential.Canonical using (Exponential)

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

module Monad.Instance.Delay.Retract {o ℓ e} {C : Category o ℓ e}
    (distributive : Distributive C)
    (DM : DelayM (Distributive.cocartesian distributive))
    (pnno : ParametrizedNNO C (Distributive.cartesian distributive))
    (_^ℕ : ∀ X → Exponential (Distributive.cartesian distributive) (ParametrizedNNO.N pnno) X) where
    
    open Category C
    open import Category.Distributive.Helper distributive
    open DelayM DM
    open HomReasoning
    open Equiv
    open Coit
    open M C
    open MR C

    open ParametrizedNNO pnno using (z; s; universal; uniform; universal-cong; commute₁; commute₂) renaming (unique to pnno-unique)
    open import Object.NaturalNumbers.Parametrized.Primitive cartesian pnno

    natural : ∀ {A B X} {f : A ⇒ X} {g : X ⇒ X} {h : B ⇒ A} → universal f g ∘ (h ⁂ id) ≈ universal (f ∘ h) g
    natural {A} {X} {Y} {f} {g} {h} = pnno-unique
      (pushˡ commute₁ ○ pushʳ (⟨⟩∘  ○ ⟨⟩-cong₂ id-comm-sym (pullʳ !-unique₂) ○ sym first∘⟨⟩))
      (extendʳ commute₂ ○ pushʳ (sym first↔second))

    -- D X is a retract of the function space (X+1)^ℕ.
    --
    -- D-to-stream samples a Delay value into a stream over ℕ: index n is
    -- i₁ x when the computation has halted with value x within n steps, and
    -- i₂ ⋆ (the ⊤ summand) while it is still running. stream-to-D
    -- coiterates a stream back into D X. D-retract packages these as a
    -- section/retraction pair, proving stream-to-D ∘ D-to-stream ≈ id.
    module _ {X : Obj} where
      open Exponential ((X + ⊤)^ℕ) using (λg; eval; λ-cong; β; subst) renaming (B^A to X+1^ℕ; η to η-exp)

      -- head of a stream: its value at index 0
      head : X+1^ℕ ⇒ X + ⊤
      head = eval ∘ ⟨ id , z ∘ ! ⟩

      -- tail of a stream: shift the index by one (drop the first element)
      tail : X+1^ℕ ⇒ X+1^ℕ
      tail = λg (eval ∘ (id ⁂ s))

      stream-to-D : X+1^ℕ ⇒ D₀ X
      stream-to-D = coit ((π₁ +₁ tail ∘ π₂) ∘ distributeʳ⁻¹ ∘ ⟨ head , id ⟩) 

      D-to-stream : D₀ X ⇒ X+1^ℕ
      D-to-stream = λg ((id +₁ !) ∘ universal out ([ i₁ , out ]))

      D-to-stream-head : head ∘ D-to-stream ≈ (id +₁ !) ∘ out
      D-to-stream-head = (pullʳ (⟨⟩∘ ○ ⟨⟩-cong₂ identityˡ (pullʳ !-unique₂)) ○ D-to-stream-z)
        where
         -- value of the sampling stream at index 0
         D-to-stream-z : eval ∘ ⟨ D-to-stream , z ∘ ! ⟩ ≈ (id +₁ !) ∘ out
         D-to-stream-z = begin
          eval ∘ ⟨ D-to-stream , z ∘ ! ⟩                                       ≈⟨ refl⟩∘⟨ (first∘⟨⟩ ○ ⟨⟩-congʳ identityʳ) ⟨
          eval ∘ (D-to-stream ⁂ id) ∘ ⟨ id , z ∘ ! ⟩                           ≈⟨ extendʳ β ⟩
          (id +₁ !) ∘ universal out ([ i₁ , out ]) ∘ ⟨ id , z ∘ ! ⟩            ≈⟨ refl⟩∘⟨ commute₁ ⟨
          (id +₁ !) ∘ out                                                      ∎

      -- the tail of the sampling stream is the sampling of `earlier`
      D-to-stream-tail : tail ∘ D-to-stream ≈ D-to-stream ∘ earlier
      D-to-stream-tail = subst ○ λ-cong (pullʳ second∘first ○ D-to-stream-s ○ ∘-resp-≈ʳ first∘first) ○ η-exp
       where
         -- shifting the index by one amounts to sampling `earlier` (per-`eval` form)
         D-to-stream-s : eval ∘ (D-to-stream ⁂ s) ≈ eval ∘ (D-to-stream ⁂ id) ∘ (earlier ⁂ id)
         D-to-stream-s = begin
           eval ∘ (D-to-stream ⁂ s)                                             ≈⟨ (refl⟩∘⟨ sym (⁂∘⁂ ○ ⁂-cong₂ identityʳ identityˡ)) ○ extendʳ β ⟩
           (id +₁ !) ∘ universal out ([ i₁ , out ]) ∘ (id ⁂ s)                  ≈⟨ refl⟩∘⟨ commute₂ ⟨
           (id +₁ !) ∘ [ i₁ , out ] ∘ universal out ([ i₁ , out ])              ≈⟨ refl⟩∘⟨ uniform refl ⟩
           (id +₁ !) ∘ universal ([ i₁ , out ] ∘ out) ([ i₁ , out ])            ≈⟨ refl⟩∘⟨ universal-cong (sym out∘earlier) refl ⟩
           (id +₁ !) ∘ universal (out ∘ earlier) ([ i₁ , out ])                 ≈⟨ refl⟩∘⟨ natural ⟨
           (id +₁ !) ∘ universal out ([ i₁ , out ]) ∘ (earlier ⁂ id)            ≈⟨ extendʳ β ⟨
           eval ∘ (D-to-stream ⁂ id) ∘ (earlier ⁂ id)                           ∎

      -- distribute a coproduct out of a product, then forget the second factor on
      -- the left summand / keep it on the right: i₁ x ↦ i₁ x, i₂ _ ↦ i₂
      distπ : (X + ⊤) × X+1^ℕ ⇒ X + X+1^ℕ
      distπ = (π₁ +₁ π₂) ∘ distributeʳ⁻¹

      distπ-i₁ : distπ ∘ (i₁ ⁂ id) ≈ i₁ ∘ π₁
      distπ-i₁ = pullʳ distributeʳ⁻¹-i₁ ○ +₁∘i₁

      distπ-i₂ : distπ ∘ (i₂ ⁂ id) ≈ i₂ ∘ π₂
      distπ-i₂ = pullʳ distributeʳ⁻¹-i₂ ○ +₁∘i₂

      D-retract : Retract (D₀ X) X+1^ℕ
      D-retract .M.Retract.section = D-to-stream
      D-retract .M.Retract.retract = stream-to-D
      D-retract .M.Retract.is-retract = coit-unique' out (stream-to-D ∘ D-to-stream) id retract-coind (id-comm ○ ∘-resp-≈ˡ (sym id+₁id))
        where
          -- coinduction hypothesis fed to coit-unique': stream-to-D ∘ D-to-stream
          -- is an out-coalgebra morphism over the identity, hence equals id
          retract-coind : out ∘ stream-to-D ∘ D-to-stream ≈ (id +₁ stream-to-D ∘ D-to-stream) ∘ out
          retract-coind = begin
            out ∘ stream-to-D ∘ D-to-stream                                                                                   ≈⟨ pullˡ (coit-commutes _) ○ pullʳ assoc ⟩ 
            (id +₁ stream-to-D) ∘ (π₁ +₁ tail ∘ π₂) ∘ (distributeʳ⁻¹ ∘ ⟨ head , id ⟩) ∘ D-to-stream                           ≈⟨ refl⟩∘⟨ refl⟩∘⟨ pullʳ (⟨⟩∘ ○ ⟨⟩-cong₂ D-to-stream-head identityˡ) ⟩
            (id +₁ stream-to-D) ∘ (π₁ +₁ tail ∘ π₂) ∘ (distributeʳ⁻¹ ∘ ⟨ (id +₁ !) ∘ out , D-to-stream ⟩)                     ≈⟨ refl⟩∘⟨ pushˡ (sym (+₁∘+₁ ○ +₁-cong₂ (π₁∘⁂ ○ identityˡ) π₂∘⁂)) ⟩ 
            (id +₁ stream-to-D) ∘ (π₁ +₁ π₂) ∘ (id ⁂ tail +₁ id ⁂ tail) ∘ distributeʳ⁻¹ ∘ ⟨ (id +₁ !) ∘ out , D-to-stream ⟩   ≈⟨ refl⟩∘⟨ refl⟩∘⟨ pullˡ (distributeʳ⁻¹-natural (tail) id id) ⟩   
            (id +₁ stream-to-D) ∘ (π₁ +₁ π₂) ∘ (distributeʳ⁻¹ ∘ ((id +₁ id) ⁂ tail)) ∘ ⟨ (id +₁ !) ∘ out , D-to-stream ⟩      ≈⟨ refl⟩∘⟨ pushʳ (pullʳ (⁂∘⟨⟩ ○ ⟨⟩-congʳ (elimˡ id+₁id))) ⟩ 
            (id +₁ stream-to-D) ∘ distπ ∘ ⟨ (id +₁ !) ∘ out , tail ∘ D-to-stream ⟩                                            ≈⟨ refl⟩∘⟨ refl⟩∘⟨ ⟨⟩-congˡ D-to-stream-tail ⟩
            (id +₁ stream-to-D) ∘ distπ ∘ ⟨ (id +₁ !) ∘ out , D-to-stream ∘ earlier ⟩                                         ≈⟨ refl⟩∘⟨ refl⟩∘⟨ (⟨⟩∘ ○ ⟨⟩-cong₂ coind-fst coind-snd) ⟨
            (id +₁ stream-to-D) ∘ distπ ∘ ⟨ (id +₁ !) , [ D-to-stream ∘ now , id ] ⟩ ∘ (id +₁ D-to-stream) ∘ out              ≈⟨ refl⟩∘⟨ (sym-assoc ○ elimˡ (sym ([]-unique coind-i₁ coind-i₂) ○ +-η)) ⟩
            (id +₁ stream-to-D) ∘ (id +₁ D-to-stream) ∘ out                                                                   ≈⟨ pullˡ +-second∘+-second ⟩ 
            (id +₁ stream-to-D ∘ D-to-stream) ∘ out                                                                           ∎
            where
              coind-fst : (id +₁ !) ∘ (id +₁ D-to-stream) ∘ out ≈ (id +₁ !) ∘ out
              coind-fst = pullˡ (+-second∘+-second ○ +₁-cong₂ refl !-unique₂)

              coind-snd : [ D-to-stream ∘ now , id ] ∘ (id +₁ D-to-stream) ∘ out ≈ D-to-stream ∘ earlier
              coind-snd = pullˡ ([]∘+-second ○ []-cong₂ refl identityˡ) ○ sym (pullˡ (∘[] ○ []-cong₂ refl identityʳ))

              -- the two injection components agree with i₁/i₂, so (via []-unique)
              -- the bracketed map distπ ∘ ⟨ ... ⟩ is the identity
              coind-i₁ : (distπ ∘ ⟨ (id +₁ !) , [ D-to-stream ∘ now , id ] ⟩) ∘ i₁ ≈ i₁
              coind-i₁ = begin
                (distπ ∘ ⟨ (id +₁ !) , [ D-to-stream ∘ now , id ] ⟩) ∘ i₁     ≈⟨ pullʳ (⟨⟩∘ ○ ⟨⟩-cong₂ (+₁∘i₁ ○ identityʳ) inject₁) ⟩
                distπ ∘ ⟨ i₁ , D-to-stream ∘ now ⟩                            ≈⟨ refl⟩∘⟨ sym (⁂∘⟨⟩ ○ ⟨⟩-cong₂ identityʳ identityˡ) ⟩
                distπ ∘ (i₁ ⁂ id) ∘ ⟨ id , D-to-stream ∘ now ⟩                ≈⟨ pullˡ distπ-i₁ ⟩
                (i₁ ∘ π₁) ∘ ⟨ id , D-to-stream ∘ now ⟩                        ≈⟨ cancelʳ project₁ ⟩
                i₁                                                            ∎

              coind-i₂ : (distπ ∘ ⟨ (id +₁ !) , [ D-to-stream ∘ now , id ] ⟩) ∘ i₂ ≈ i₂
              coind-i₂ = begin
                (distπ ∘ ⟨ (id +₁ !) , [ D-to-stream ∘ now , id ] ⟩) ∘ i₂     ≈⟨ pullʳ (⟨⟩∘ ○ ⟨⟩-cong₂ +₁∘i₂ inject₂) ⟩
                distπ ∘ ⟨ i₂ ∘ ! , id ⟩                                       ≈⟨ refl⟩∘⟨ sym (⁂∘⟨⟩ ○ ⟨⟩-congˡ identityˡ) ⟩
                distπ ∘ (i₂ ⁂ id) ∘ ⟨ ! , id ⟩                                ≈⟨ pullˡ distπ-i₂ ⟩
                (i₂ ∘ π₂) ∘ ⟨ ! , id ⟩                                        ≈⟨ cancelʳ project₂ ⟩
                i₂                                                            ∎

    -- D-to-stream is natural
    module _ {X Y : Obj} (f : X ⇒ Y) where
      private
        module Xᵉ = Exponential ((X + ⊤)^ℕ)
        module Yᵉ = Exponential ((Y + ⊤)^ℕ)
        
      open D-Monad

      -- functorial action of (- + 1)^ℕ on f
      +1^ℕ : Xᵉ.B^A ⇒ Yᵉ.B^A
      +1^ℕ = Yᵉ.λg ((f +₁ !) ∘ Xᵉ.eval)

      -- naturality of the section D-to-stream
      D-to-stream-natural : +1^ℕ ∘ D-to-stream ≈ D-to-stream ∘ D₁ f
      D-to-stream-natural = begin
        +1^ℕ ∘ D-to-stream                                                 ≈⟨ Yᵉ.subst ⟩
        Yᵉ.λg (((f +₁ !) ∘ Xᵉ.eval) ∘ (D-to-stream ⁂ id))                   ≈⟨ Yᵉ.λ-cong lhs ⟩
        Yᵉ.λg ((f +₁ !) ∘ universal out ([ i₁ , out ]))                     ≈⟨ Yᵉ.λ-cong rhs ⟨
        Yᵉ.λg (((id +₁ !) ∘ universal out ([ i₁ , out ])) ∘ (D₁ f ⁂ id))    ≈⟨ Yᵉ.subst ⟨
        D-to-stream ∘ D₁ f                                                  ∎
        where
          -- the sampling map universal out [i₁,out] is natural in f
          g-nat : universal out ([ i₁ , out ]) ∘ (D₁ f ⁂ id) ≈ (f +₁ D₁ f) ∘ universal out ([ i₁ , out ])
          g-nat = begin
            universal out ([ i₁ , out ]) ∘ (D₁ f ⁂ id)        ≈⟨ natural ⟩
            universal (out ∘ D₁ f) ([ i₁ , out ])             ≈⟨ universal-cong (D₁-commutes f) refl ⟩
            universal ((f +₁ D₁ f) ∘ out) ([ i₁ , out ])      ≈⟨ uniform (sym ([]∘+₁ ○ []-unique (pullʳ inject₁ ○ inject₁) (pullʳ inject₂ ○ sym (D₁-commutes f)))) ⟨
            (f +₁ D₁ f) ∘ universal out ([ i₁ , out ])        ∎

          lhs : ((f +₁ !) ∘ Xᵉ.eval) ∘ (D-to-stream ⁂ id) ≈ (f +₁ !) ∘ universal out ([ i₁ , out ])
          lhs = assoc ○ (refl⟩∘⟨ Xᵉ.β) ○ pullˡ (+₁∘+₁ ○ +₁-cong₂ identityʳ !-unique₂)

          rhs : ((id +₁ !) ∘ universal out ([ i₁ , out ])) ∘ (D₁ f ⁂ id) ≈ (f +₁ !) ∘ universal out ([ i₁ , out ])
          rhs = assoc ○ (refl⟩∘⟨ g-nat) ○ pullˡ (+₁∘+₁ ○ +₁-cong₂ identityˡ !-unique₂)

      stream-to-D-natural : D₁ f ∘ stream-to-D ≈ stream-to-D ∘ +1^ℕ
      stream-to-D-natural = coit-unique' c (D₁ f ∘ stream-to-D) (stream-to-D ∘ +1^ℕ) coalg₁ coalg₂
        where
          -- the stream coalgebra of stream-to-D with the output value post-composed with f
          c : Xᵉ.B^A ⇒ Y + Xᵉ.B^A
          c = (f ∘ π₁ +₁ tail ∘ π₂) ∘ distributeʳ⁻¹ ∘ ⟨ head , id ⟩

          cˣ : Xᵉ.B^A ⇒ X + Xᵉ.B^A
          cˣ = (π₁ +₁ tail ∘ π₂) ∘ distributeʳ⁻¹ ∘ ⟨ head , id ⟩

          cʸ : Yᵉ.B^A ⇒ Y + Yᵉ.B^A
          cʸ = (π₁ +₁ tail ∘ π₂) ∘ distributeʳ⁻¹ ∘ ⟨ head , id ⟩

          head-nat : head ∘ +1^ℕ ≈ (f +₁ !) ∘ head
          head-nat = begin
            head ∘ +1^ℕ                                ≈⟨ pullʳ (⟨⟩∘ ○ ⟨⟩-cong₂ identityˡ (pullʳ !-unique₂)) ⟩
            Yᵉ.eval ∘ ⟨ +1^ℕ , z ∘ ! ⟩                 ≈⟨ refl⟩∘⟨ (⁂∘⟨⟩ ○ ⟨⟩-cong₂ identityʳ identityˡ) ⟨
            Yᵉ.eval ∘ (+1^ℕ ⁂ id) ∘ ⟨ id , z ∘ ! ⟩     ≈⟨ pullˡ Yᵉ.β ○ assoc ⟩  
            (f +₁ !) ∘ head                            ∎

          eq : (Yᵉ.eval ∘ (id ⁂ s)) ∘ (+1^ℕ ⁂ id) ≈ (f +₁ !) ∘ Xᵉ.eval ∘ (id ⁂ s)
          eq = begin
            (Yᵉ.eval ∘ (id ⁂ s)) ∘ (+1^ℕ ⁂ id)         ≈⟨ pullʳ (⁂∘⁂ ○ ⁂-cong₂ identityˡ identityʳ) ⟩
            Yᵉ.eval ∘ (+1^ℕ ⁂ s)                       ≈⟨ refl⟩∘⟨ (⁂-cong₂ (sym identityʳ) (sym identityˡ) ○ sym ⁂∘⁂) ⟩
            Yᵉ.eval ∘ (+1^ℕ ⁂ id) ∘ (id ⁂ s)           ≈⟨ extendʳ Yᵉ.β ⟩
            (f +₁ !) ∘ Xᵉ.eval ∘ (id ⁂ s)              ∎

          tail-nat : tail ∘ +1^ℕ ≈ +1^ℕ ∘ tail
          tail-nat = begin
            tail ∘ +1^ℕ                                   ≈⟨ Yᵉ.subst ⟩
            Yᵉ.λg ((Yᵉ.eval ∘ (id ⁂ s)) ∘ (+1^ℕ ⁂ id))     ≈⟨ Yᵉ.λ-cong eq ⟩
            Yᵉ.λg ((f +₁ !) ∘ Xᵉ.eval ∘ (id ⁂ s))          ≈⟨ Yᵉ.λ-cong (pullʳ Xᵉ.β) ⟨
            Yᵉ.λg (((f +₁ !) ∘ Xᵉ.eval) ∘ (tail ⁂ id))     ≈⟨ Yᵉ.subst ⟨
            +1^ℕ ∘ tail                                   ∎

          coalg-nat : cʸ ∘ +1^ℕ ≈ (id +₁ +1^ℕ) ∘ c
          coalg-nat = begin
            cʸ ∘ +1^ℕ                                                                       ≈⟨ pullʳ (pullʳ (⟨⟩∘ ○ ⟨⟩-congˡ identityˡ )) ⟩
            (π₁ +₁ tail ∘ π₂) ∘ distributeʳ⁻¹ ∘ ⟨ head ∘ +1^ℕ , +1^ℕ ⟩                      ≈⟨ refl⟩∘⟨ refl⟩∘⟨ ⟨⟩-congʳ head-nat ⟩
            (π₁ +₁ tail ∘ π₂) ∘ distributeʳ⁻¹ ∘ ⟨ (f +₁ !) ∘ head , +1^ℕ ⟩                  ≈⟨ refl⟩∘⟨ refl⟩∘⟨ (⟨⟩∘ ○ ⟨⟩-cong₂ (pullʳ project₁) (cancelʳ project₂)) ⟨ 
            (π₁ +₁ tail ∘ π₂) ∘ distributeʳ⁻¹ ∘ ((f +₁ !) ⁂ +1^ℕ) ∘ ⟨ head , id ⟩           ≈⟨ refl⟩∘⟨ extendʳ (distributeʳ⁻¹-natural +1^ℕ f !) ⟨ 
            (π₁ +₁ tail ∘ π₂) ∘ (f  ⁂ +1^ℕ +₁ ! ⁂ +1^ℕ) ∘  distributeʳ⁻¹ ∘ ⟨ head , id ⟩    ≈⟨ extendʳ (+₁∘+₁ ○ +₁-cong₂ (project₁ ○ sym identityˡ)
                                                                                                                          (extendˡ project₂ ○ pushˡ tail-nat) ○ sym +₁∘+₁) ⟩
            (id +₁ +1^ℕ) ∘ c                                                                ∎


          -- D₁ f ∘ stream-to-D is an out-coalgebra morphism over c
          coalg₁ : out ∘ (D₁ f ∘ stream-to-D) ≈ (id +₁ (D₁ f ∘ stream-to-D)) ∘ c
          coalg₁ = begin
            out ∘ (D₁ f ∘ stream-to-D)               ≈⟨ extendʳ (D₁-commutes f) ⟩
            (f +₁ D₁ f) ∘ out ∘ stream-to-D          ≈⟨ refl⟩∘⟨ coit-commutes _ ⟩
            (f +₁ D₁ f) ∘ (id +₁ stream-to-D) ∘ cˣ   ≈⟨ pullˡ (+₁∘+₁ ○ +₁-cong₂ identityʳ refl) ⟩
            (f +₁ D₁ f ∘ stream-to-D) ∘ cˣ           ≈⟨ extendʳ (+₁∘+₁ ○ +₁-cong₂ (sym identityˡ) refl ○ sym +₁∘+₁) ⟩
            (id +₁ D₁ f ∘ stream-to-D) ∘ c           ∎

          -- stream-to-D ∘ +1^ℕ is an out-coalgebra morphism over c
          coalg₂ : out ∘ (stream-to-D ∘ +1^ℕ) ≈ (id +₁ (stream-to-D ∘ +1^ℕ)) ∘ c
          coalg₂ = begin
            out ∘ (stream-to-D ∘ +1^ℕ)               ≈⟨ extendʳ (coit-commutes _) ⟩
            (id +₁ stream-to-D) ∘ cʸ ∘ +1^ℕ          ≈⟨ refl⟩∘⟨ coalg-nat ⟩
            (id +₁ stream-to-D) ∘ (id +₁ +1^ℕ) ∘ c   ≈⟨ pullˡ (+₁∘+₁ ○ +₁-cong₂ identityˡ refl) ⟩
            (id +₁ stream-to-D ∘ +1^ℕ) ∘ c           ∎