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.Functor hiding (id)
open import Categories.NaturalTransformation hiding (id)

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 D-Monad
    open M C
    open MR C
    open Functor

    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))

    module Stream {X : Obj} = Exponential ((X + ) ^ℕ)
    open Stream using (λg; β; eval; λ-cong; η-id; λ-unique′; subst)
      renaming (η to η-exp)

    infix 30 _+1^ℕ
    _+1^ℕ : Obj  Obj
    X +1^ℕ = Stream.B^A {X}

    Stream-Functor : Endofunctor C
    Stream-Functor .F₀ X = X +1^ℕ
    Stream-Functor .F₁ f = λg ((f +₁ !)  eval)
    Stream-Functor .identity = λ-cong
      (∘-resp-≈ˡ ((+₁-cong₂ refl !-unique₂)  id+₁id)  identityˡ)  η-id
    Stream-Functor .homomorphism {f = f} {g = g} = λ-unique′ (begin
      eval  (λg (((g  f) +₁ !)  eval) ×₁ id)                          ≈⟨ β 
      ((g  f) +₁ !)  eval                                              ≈⟨ ∘-resp-≈ˡ (+₁-cong₂ refl !-unique₂  sym +₁∘+₁) 
      ((g +₁ !)  (f +₁ !))  eval                                       ≈⟨ assoc 
      (g +₁ !)  (f +₁ !)  eval                                         ≈⟨ refl⟩∘⟨ β 
      (g +₁ !)  eval  (λg ((f +₁ !)  eval) ×₁ id)                     ≈⟨ extendʳ β 
      eval  (λg ((g +₁ !)  eval) ×₁ id)  (λg ((f +₁ !)  eval) ×₁ id) ≈⟨ refl⟩∘⟨ (×₁∘×₁  ×₁-cong₂ refl identity²) 
      eval  ((λg ((g +₁ !)  eval)  λg ((f +₁ !)  eval)) ×₁ id)       )
    Stream-Functor .F-resp-≈ eq = λ-cong
      (∘-resp-≈ˡ (+₁-cong₂ eq refl))

    module Stream-Functor = Functor Stream-Functor

    -- 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
      private
        X+1^ℕ : Obj
        X+1^ℕ = X +1^ℕ

      -- 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
      open D-Monad

      -- naturality of the section D-to-stream
      D-to-stream-natural : Stream-Functor.₁ f  D-to-stream  D-to-stream  D₁ f
      D-to-stream-natural = begin
          Stream-Functor.₁ f  D-to-stream                               ≈⟨ subst 
          λg (((f +₁ !)  eval)  (D-to-stream ×₁ id))                   ≈⟨ λ-cong lhs 
          λg ((f +₁ !)  universal out ([ i₁ , out ]))                   ≈⟨ λ-cong rhs 
          λg (((id +₁ !)  universal out ([ i₁ , out ]))  (D₁ f ×₁ id)) ≈⟨ 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 +₁ !)  eval)  (D-to-stream ×₁ id)  (f +₁ !)  universal out ([ i₁ , out ])
            lhs = assoc  (refl⟩∘⟨ β)  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  Stream-Functor.₁ f
      stream-to-D-natural = coit-unique' c (D₁ f  stream-to-D) (stream-to-D  Stream-Functor.₁ f) coalg₁ coalg₂
          where
            -- the stream coalgebra of stream-to-D with the output value post-composed with f
            c : X +1^ℕ  Y + X +1^ℕ
            c = (f  π₁ +₁ tail  π₂)  distributeʳ⁻¹   head , id 

             : X +1^ℕ  X + X +1^ℕ
             = (π₁ +₁ tail  π₂)  distributeʳ⁻¹   head , id 

             : Y +1^ℕ  Y + Y +1^ℕ
             = (π₁ +₁ tail  π₂)  distributeʳ⁻¹   head , id 

            head-nat : head  Stream-Functor.₁ f  (f +₁ !)  head
            head-nat = begin
              head  Stream-Functor.₁ f                          ≈⟨ pullʳ (⟨⟩∘  ⟨⟩-cong₂ identityˡ (pullʳ !-unique₂)) 
              eval   Stream-Functor.₁ f , z  !               ≈⟨ refl⟩∘⟨ (×₁∘⟨⟩  ⟨⟩-cong₂ identityʳ identityˡ) 
              eval  (Stream-Functor.₁ f ×₁ id)   id , z  !  ≈⟨ pullˡ β  assoc   
              (f +₁ !)  head                                    

            eq : (eval  (id ×₁ s))  (Stream-Functor.₁ f ×₁ id)  (f +₁ !)  eval  (id ×₁ s)
            eq = begin
              (eval  (id ×₁ s))  (Stream-Functor.₁ f ×₁ id) ≈⟨ pullʳ (×₁∘×₁  ×₁-cong₂ identityˡ identityʳ) 
              eval  (Stream-Functor.₁ f ×₁ s)                ≈⟨ refl⟩∘⟨ (×₁-cong₂ (sym identityʳ) (sym identityˡ)  sym ×₁∘×₁) 
              eval  (Stream-Functor.₁ f ×₁ id)  (id ×₁ s)   ≈⟨ extendʳ β 
              (f +₁ !)  eval  (id ×₁ s)                     

            tail-nat : tail  Stream-Functor.₁ f  Stream-Functor.₁ f  tail
            tail-nat = begin
              tail  Stream-Functor.₁ f                            ≈⟨ subst 
              λg ((eval  (id ×₁ s))  (Stream-Functor.₁ f ×₁ id)) ≈⟨ λ-cong eq 
              λg ((f +₁ !)  eval  (id ×₁ s))                     ≈⟨ λ-cong (pullʳ β) 
              λg (((f +₁ !)  eval)  (tail ×₁ id))                ≈⟨ subst 
              Stream-Functor.₁ f  tail                            

            coalg-nat :   Stream-Functor.₁ f  (id +₁ Stream-Functor.₁ f)  c
            coalg-nat = begin
                Stream-Functor.₁ f                                                                                    ≈⟨ pullʳ (pullʳ (⟨⟩∘  ⟨⟩-congˡ identityˡ )) 
              (π₁ +₁ tail  π₂)  distributeʳ⁻¹   head  Stream-Functor.₁ f , Stream-Functor.₁ f                      ≈⟨ refl⟩∘⟨ refl⟩∘⟨ ⟨⟩-congʳ head-nat 
              (π₁ +₁ tail  π₂)  distributeʳ⁻¹   (f +₁ !)  head , Stream-Functor.₁ f                                ≈⟨ refl⟩∘⟨ refl⟩∘⟨ (⟨⟩∘  ⟨⟩-cong₂ (pullʳ project₁) (cancelʳ project₂))  
              (π₁ +₁ tail  π₂)  distributeʳ⁻¹  ((f +₁ !) ×₁ Stream-Functor.₁ f)   head , id                        ≈⟨ refl⟩∘⟨ extendʳ (distributeʳ⁻¹-natural (Stream-Functor.₁ f) f !)  
              (π₁ +₁ tail  π₂)  (f  ×₁ Stream-Functor.₁ f +₁ ! ×₁ Stream-Functor.₁ f)   distributeʳ⁻¹   head , id  ≈⟨ extendʳ (+₁∘+₁  +₁-cong₂ (project₁  sym identityˡ)
                                                                                                                            (extendˡ project₂  pushˡ tail-nat)  sym +₁∘+₁) 
              (id +₁ Stream-Functor.₁ f)  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)   ≈⟨ pullˡ (+₁∘+₁  +₁-cong₂ identityʳ refl) 
              (f +₁ D₁ f  stream-to-D)           ≈⟨ extendʳ (+₁∘+₁  +₁-cong₂ (sym identityˡ) refl  sym +₁∘+₁) 
              (id +₁ D₁ f  stream-to-D)  c         

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

    D⇒Stream : NaturalTransformation D.F Stream-Functor
    D⇒Stream = ntHelper record
      { η = λ _  D-to-stream
      ; commute = λ f  sym (D-to-stream-natural f)
      }

    Stream⇒D : NaturalTransformation Stream-Functor D.F
    Stream⇒D = ntHelper record
      { η = λ _  stream-to-D
      ; commute = λ f  sym (stream-to-D-natural f)
      }

    module D⇒Stream = NaturalTransformation D⇒Stream
    module Stream⇒D = NaturalTransformation Stream⇒D

    Stream-D-retract :  X  Stream⇒D.η X  D⇒Stream.η X  id
    Stream-D-retract X = M.Retract.is-retract D-retract