open import Categories.Category.Core
open import Categories.Object.Product.Core using (Product)
open import Categories.Object.Terminal
open import Data.Product using (_,_)
open import Categories.Category.Distributive
open import Monad.Instance.Delay
open import Categories.Monad hiding (id)
open import Categories.Monad.Relative renaming (Monad to RMonad)
open import Categories.Functor.Core
open import Categories.Monad.Strong

open import Monad.Helper

import Categories.Morphism as Mor
import Categories.Morphism.Reasoning as MR
import Categories.Morphism.Properties as MP

module Algebra.Search.Retraction {o ℓ e} {C : Category o ℓ e} (distributive : Distributive C) (D : DelayM (Distributive.cocartesian distributive)) where
  open Category C
  open import Category.Distributive.Helper distributive
  open Bundles
  open import Monad.Instance.Delay.Commutative distributive D
  open import Monad.Instance.Delay.Zip distributive D
  open τ-mod
  open σ-mod
  open HomReasoning
  open Equiv
  open Mor C
  open MR C
  open MP C
  open DelayM D
  open import Monad.Instance.Delay.Guarded cocartesian D
  open D-Kleisli
  open D-Monad
  open Later∘Extend
  open import Algebra.Search cocartesian D


  tau-retract : ∀ X (algebra : Search-Algebra) → Retract (D₀ ((Search-Algebra.A algebra) × X)) (Search-Algebra.A algebra × D₀ X)
  tau-retract X algebra = record 
    { section = ⟨ α ∘ D₁ π₁ , D₁ π₂ ⟩
    ; retract = τ 
    ; is-retract = begin 
        τ ∘ ⟨ α ∘ D₁ π₁ , D₁ π₂ ⟩                                                            ≈⟨ introˡ product-is-retract ⟩
        ((extend [ now , [ τ , σ ] ] ∘ zip) ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩) ∘ τ ∘ ⟨ α ∘ D₁ π₁ , D₁ π₂ ⟩ ≈⟨ pullʳ epi-helper ⟩
        (extend [ now , [ τ , σ ] ] ∘ zip) ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩                               ≈⟨ product-is-retract ⟩
        id                                                                                   ∎ 
    }
    
    where
    open Search-Algebra algebra using (A; α) renaming (identity₁ to s-identity₁; identity₂ to s-identity₂)
    open Retract (product-retract A X) using () renaming (is-retract to product-is-retract)
    top-square : ((now ×₁ id) +₁ ⟨ later ∘ D₁ π₁ , D₁ π₂ ⟩) ∘ out ≈ distributeˡ⁻¹ ∘ (id ×₁ out) ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩
    top-square = out⁻¹-epi (((now ×₁ id) +₁ ⟨ later ∘ D₁ π₁ , D₁ π₂ ⟩) ∘ out) (distributeˡ⁻¹ ∘ (id ×₁ out) ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩) (begin 
      (((now ×₁ id) +₁ ⟨ later ∘ D₁ π₁ , D₁ π₂ ⟩) ∘ out) ∘ out⁻¹ ≈⟨ cancelʳ out∘out⁻¹ ⟩ 
      (now ×₁ id) +₁ ⟨ later ∘ D₁ π₁ , D₁ π₂ ⟩                   ≈⟨ []-unique unique₁ unique₂ ⟩ 
      (distributeˡ⁻¹ ∘ (id ×₁ out) ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩) ∘ out⁻¹  ∎)
      where
      unique₁ : ((distributeˡ⁻¹ ∘ (id ×₁ out) ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩) ∘ out⁻¹) ∘ i₁ ≈ i₁ ∘ (now ×₁ id)
      unique₁ = begin 
        ((distributeˡ⁻¹ ∘ (id ×₁ out) ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩) ∘ out⁻¹) ∘ i₁ ≈⟨ assoc ⟩ 
        (distributeˡ⁻¹ ∘ (id ×₁ out) ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩) ∘ now          ≈⟨ pullʳ (pullʳ ⟨⟩∘) ⟩ 
        distributeˡ⁻¹ ∘ (id ×₁ out) ∘ ⟨ D₁ π₁ ∘ now , D₁ π₂ ∘ now ⟩      ≈⟨ refl⟩∘⟨ refl⟩∘⟨ ⟨⟩-cong₂ DK.identityʳ DK.identityʳ ⟩ 
        distributeˡ⁻¹ ∘ (id ×₁ out) ∘ ⟨ now ∘ π₁ , now ∘ π₂ ⟩            ≈⟨ refl⟩∘⟨ (×₁∘⟨⟩ ○ ⟨⟩-cong₂ identityˡ (pullˡ unitlaw)) ⟩ 
        distributeˡ⁻¹ ∘ ⟨ now ∘ π₁ , i₁ ∘ π₂ ⟩                           ≈˘⟨ refl⟩∘⟨ (×₁∘⟨⟩ ○ ⟨⟩-cong₂ identityˡ refl) ⟩ 
        distributeˡ⁻¹ ∘ (id ×₁ i₁) ∘ ⟨ now ∘ π₁ , π₂ ⟩                   ≈⟨ pullˡ distributeˡ⁻¹-i₁ ⟩ 
        i₁ ∘ ⟨ now ∘ π₁ , π₂ ⟩                                           ≈⟨ refl⟩∘⟨ (⟨⟩-cong₂ refl (sym identityˡ)) ⟩ 
        i₁ ∘ (now ×₁ id)                                                 ∎
      unique₂ : ((distributeˡ⁻¹ ∘ (id ×₁ out) ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩) ∘ out⁻¹) ∘ i₂ ≈ i₂ ∘ ⟨ later ∘ D₁ π₁ , D₁ π₂ ⟩
      unique₂ = begin 
        ((distributeˡ⁻¹ ∘ (id ×₁ out) ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩) ∘ out⁻¹) ∘ i₂ ≈⟨ assoc ⟩ 
        (distributeˡ⁻¹ ∘ (id ×₁ out) ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩) ∘ later        ≈⟨ pullʳ (pullʳ ⟨⟩∘) ⟩ 
        distributeˡ⁻¹ ∘ (id ×₁ out) ∘ ⟨ D₁ π₁ ∘ later , D₁ π₂ ∘ later ⟩  ≈⟨ refl⟩∘⟨ refl⟩∘⟨ ⟨⟩-cong₂ (sym (later-extend-comm (now ∘ π₁))) (sym (later-extend-comm (now ∘ π₂))) ⟩ 
        distributeˡ⁻¹ ∘ (id ×₁ out) ∘ ⟨ later ∘ D₁ π₁ , later ∘ D₁ π₂ ⟩  ≈⟨ refl⟩∘⟨ (×₁∘⟨⟩ ○ ⟨⟩-cong₂ identityˡ (pullˡ laterlaw)) ⟩ 
        distributeˡ⁻¹ ∘ ⟨ later ∘ D₁ π₁ , i₂ ∘ D₁ π₂ ⟩                   ≈˘⟨ refl⟩∘⟨ (×₁∘⟨⟩ ○ ⟨⟩-cong₂ identityˡ refl) ⟩ 
        distributeˡ⁻¹ ∘ (id ×₁ i₂) ∘ ⟨ later ∘ D₁ π₁ , D₁ π₂ ⟩           ≈⟨ pullˡ distributeˡ⁻¹-i₂ ⟩ 
        i₂ ∘ ⟨ later ∘ D₁ π₁ , D₁ π₂ ⟩                                  ∎
    eq : D₁ π₁ ≈ D₁ α ∘ D₁ π₁ ∘ τ ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩
    eq = extend-unique (now ∘ π₁) (D₁ α ∘ D₁ π₁ ∘ τ ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩) (begin 
      out ∘ D₁ α ∘ D₁ π₁ ∘ τ ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩                                                  ≈⟨ extendʳ (D₁-commutes α) ⟩ 
      (α +₁ D₁ α) ∘ out ∘ D₁ π₁ ∘ τ ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩                                           ≈⟨ refl⟩∘⟨ (extendʳ (D₁-commutes π₁)) ⟩ 
      (α +₁ D₁ α) ∘ (π₁ +₁ D₁ π₁) ∘ out ∘ τ ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩                                   ≈⟨ refl⟩∘⟨ (refl⟩∘⟨ (extendʳ τ-commutes)) ⟩ 
      (α +₁ D₁ α) ∘ (π₁ +₁ D₁ π₁) ∘ (id +₁ τ) ∘ (distributeˡ⁻¹ ∘ (id ×₁ out)) ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩ ≈⟨ refl⟩∘⟨ refl⟩∘⟨ refl⟩∘⟨ (assoc ○ sym top-square) ⟩ 
      (α +₁ D₁ α) ∘ (π₁ +₁ D₁ π₁) ∘ (id +₁ τ) ∘ ((now ×₁ id) +₁ ⟨ later ∘ D₁ π₁ , D₁ π₂ ⟩) ∘ out  ≈⟨ pullˡ +₁∘+₁ ○ pullˡ +₁∘+₁ ○ pullˡ +₁∘+₁ ⟩ 
      (((α ∘ π₁) ∘ id) ∘ (now ×₁ id) +₁ ((D₁ α ∘ D₁ π₁) ∘ τ) ∘ ⟨ later ∘ D₁ π₁ , D₁ π₂ ⟩) ∘ out   ≈⟨ (+₁-cong₂ (pullʳ identityˡ ○ pullʳ project₁) assoc²αε) ⟩∘⟨refl ⟩ 
      (α ∘ now ∘ π₁ +₁ D₁ α ∘ D₁ π₁ ∘ τ ∘ ⟨ later ∘ D₁ π₁ , D₁ π₂ ⟩) ∘ out                        ≈⟨ +₁-cong₂ (cancelˡ s-identity₁) (sym (refl⟩∘⟨ refl⟩∘⟨ refl⟩∘⟨ (×₁∘⟨⟩ ○ ⟨⟩-cong₂ refl identityˡ))) ⟩∘⟨refl ⟩ 
      (π₁ +₁ D₁ α ∘ D₁ π₁ ∘ τ ∘ (later ×₁ id) ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩) ∘ out                          ≈⟨ (+₁-cong₂ refl (refl⟩∘⟨ refl⟩∘⟨ extendʳ (refl⟩∘⟨ (×₁-cong₂ refl (sym D.F.identity)) ○ τ-natural later id))) ⟩∘⟨refl ⟩
      (π₁ +₁ D₁ α ∘ D₁ π₁ ∘ D₁ (later ×₁ id) ∘ τ ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩) ∘ out                       ≈⟨ (+₁-cong₂ refl (refl⟩∘⟨ (extendʳ (sym D.F.homomorphism ○ D.F.F-resp-≈ project₁ ○ D.F.homomorphism)))) ⟩∘⟨refl ⟩ 
      (π₁ +₁ D₁ α ∘ D₁ later ∘ D₁ π₁ ∘ τ ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩) ∘ out                               ≈⟨ ([]-cong₂ (pushˡ (sym unitlaw)) (refl⟩∘⟨ (pullˡ (sym D.F.homomorphism ○ D.F.F-resp-≈ s-identity₂)))) ⟩∘⟨refl ⟩ 
      [ out ∘ now ∘ π₁ , i₂ ∘ D₁ α ∘ D₁ π₁ ∘ τ ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩ ] ∘ out                        ∎)
    epi-helper : ⟨ D₁ π₁ , D₁ π₂ ⟩ ∘ τ ∘ ⟨ α ∘ D₁ π₁ , D₁ π₂ ⟩ ≈ ⟨ D₁ π₁ , D₁ π₂ ⟩
    epi-helper = begin 
      ⟨ D₁ π₁ , D₁ π₂ ⟩ ∘ τ ∘ ⟨ α ∘ D₁ π₁ , D₁ π₂ ⟩                              ≈⟨ ⟨⟩∘ ⟩ 
      ⟨ D₁ π₁ ∘ τ ∘ ⟨ α ∘ D₁ π₁ , D₁ π₂ ⟩ , D₁ π₂ ∘ τ ∘ ⟨ α ∘ D₁ π₁ , D₁ π₂ ⟩ ⟩  ≈⟨ ⟨⟩-cong₂ (sym (refl⟩∘⟨ refl⟩∘⟨ (×₁∘⟨⟩ ○ ⟨⟩-cong₂ refl identityˡ))) (pullˡ τ-identityˡ) ⟩
      ⟨ D₁ π₁ ∘ τ ∘ (α ×₁ id) ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩ , π₂ ∘ ⟨ α ∘ D₁ π₁ , D₁ π₂ ⟩ ⟩ ≈⟨ ⟨⟩-cong₂ (refl⟩∘⟨ (extendʳ (refl⟩∘⟨ (×₁-cong₂ refl (sym D.F.identity)) ○ τ-natural α id))) project₂ ⟩ 
      ⟨ D₁ π₁ ∘ D₁ (α ×₁ id) ∘ τ ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩ , D₁ π₂ ⟩                   ≈⟨ ⟨⟩-cong₂ (extendʳ (sym D.F.homomorphism ○ D.F.F-resp-≈ project₁ ○ D.F.homomorphism)) refl ⟩
      ⟨ D₁ α ∘ D₁ π₁ ∘ τ ∘ ⟨ D₁ π₁ , D₁ π₂ ⟩ , D₁ π₂ ⟩                           ≈⟨ ⟨⟩-cong₂ (sym eq) refl ⟩
      ⟨ D₁ π₁ , D₁ π₂ ⟩                                                          ∎