Derivative.Indexed.Mu

-- TODO: Clean up imports
{-# OPTIONS -WnoUnsupportedIndexedMatch --safe #-}
module Derivative.Indexed.Mu where

open import Derivative.Indexed.Container
open import Derivative.Indexed.Derivative
open import Derivative.Indexed.Univalence
open import Derivative.Indexed.ChainRule

open import Derivative.Prelude
open import Derivative.Basics.Decidable
open import Derivative.Basics.Equiv
open import Derivative.Basics.Maybe
open import Derivative.Basics.Sigma
open import Derivative.Basics.Sum
open import Derivative.Basics.W

open import Derivative.Isolated
open import Derivative.Remove

open import Cubical.Foundations.Path
open import Cubical.Foundations.Equiv.Properties using (congEquiv)
open import Cubical.Foundations.Transport using (substEquiv)
open import Cubical.Data.Sigma
open import Cubical.Functions.FunExtEquiv
open import Cubical.Functions.Embedding
open import Cubical.Functions.Surjection

private
  variable
     : Level
    Ix : Type

open Container

μ : (F : Container  (Maybe Ix))  Container  Ix
μ {} {Ix} F = shape  pos module μ where
  open Container F renaming (Shape to S ; Pos to P)

  shape : Type 
  shape = W S (P nothing)

  pos : Ix  shape  Type 
  pos ix = Wᴰ S (P nothing) (P (just ix))

module _ (F : Container  (Maybe Ix)) where
  open Container F renaming (Shape to S ; Pos to P)

  μ-in-equiv : Equiv (F [ μ F ]) (μ F)
  μ-in-equiv = [ shape ◁≃ pos ] where
    shape : Shape (F [ μ F ])  Shape (μ F)
    shape = W-in-equiv

    pos :  ix w*  Pos (μ F) ix (W-in w*)  Pos (F [ μ F ]) ix w*
    pos ix (s , ws) = Wᴰ-out-equiv S (P nothing) (P (just ix)) s ws

  μ-in : F [ μ F ]  μ F
  μ-in = Equiv.as-⊸ μ-in-equiv

  μ-out-equiv : Equiv (μ F) (F [ μ F ])
  μ-out-equiv = [ shape ◁≃ pos ] where
    shape : Shape (μ F)  Shape (F [ μ F ])
    shape = W-out-equiv

    pos :  ix s*  Pos (F [ μ F ]) ix (W-out s*)  Pos (μ F) ix s*
    pos ix (sup s f) = Wᴰ-in-equiv _ _ _ s f
    {-# INLINE pos #-}

  μ-out : μ F  F [ μ F ]
  μ-out = Equiv.as-⊸ μ-out-equiv

  μ-in-out : μ-in  μ-out  id (F [ μ F ])
  μ-in-out = ⊸≡-ext shape-path λ ix (s , f)  funExt (pos-path ix s f) where
    shape-path : W-in  W-out  idfun _
    shape-path = refl

    pos-path :  ix s
       (f : P nothing s  W S (P _))
       (p : (F [ μ F ]) .Pos ix _)  Wᴰ-out S _ _ s f (Wᴰ-in S _ _ s f p)  p
    pos-path ix s f (inl _) = refl
    pos-path ix s f (inr _) = refl

  μ-rec-Π : (G : Container  Ix)
     (F [ G ]) Π⊸ G
     μ F Π⊸ G
  μ-rec-Π G α = W-elim goal where
    module _ (s : S) (f : P nothing s  W S (P nothing)) where
      goal : (∀ p  Σ[ t  G .Shape ]  ix  Pos G ix t  Pos (μ F) ix (f p))
         Σ[ t  G .Shape ]  ix  Pos G ix t  Pos (μ F) ix (sup s f)
      goal rec .fst = α (s , fst  rec) .fst
      goal rec .snd = λ ix  α (s , fst  rec) .snd ix ∙ₑ ⊎-right-≃ (Σ-cong-equiv-snd λ p  rec p .snd ix) ∙ₑ Wᴰ-in-equiv _ _ _ s f

  μ-rec : (G : Container  Ix)
     (F [ G ])  G
     μ F  G
  μ-rec G φ = [ shape  pos ] module μ-rec where
    open Container G renaming (Shape to T ; Pos to Q)

    module φ = _⊸_ φ

    shape : W S (P nothing)  T
    shape = W-rec φ.shape

    shape-β : shape  φ.shape  Σ-map-snd  s (f : P  s  W S _)  shape  f)  W-out
    shape-β = funExt λ { (sup s f)  refl }

    pos-fun :  i  (s : W S (P ))  Q i (shape s)  Pos (μ F) i s
    pos-fun i (sup s f) =
      Q i (φ.shape (s , shape  f))
        →⟨ equivFun $ φ.pos i (s , shape  f) 
      (P (just i) s  (Σ[ p  P nothing s ] Q i (shape (f p))))
        →⟨ ⊎-map-right (Σ-map-snd $ pos-fun i  f) 
      (P (just i) s  (Σ[ p  P nothing s ] Pos (μ F) i (f p)))
        →⟨ Wᴰ-in S (P nothing) (P (just i)) s f 
      Wᴰ S (P nothing) (P (just i)) (sup s f)
        →∎

    is-equiv-pos-fun :  i (w : W S (P ))  isEquiv (pos-fun i w)
    is-equiv-pos-fun i (sup s f) = goal where
      step-1 : isEquiv (⊎-map-right {A = P (just i) s} (Σ-map-snd $ pos-fun i  f))
      step-1 = isEquiv→isEquiv-⊎-map-right $ isEquiv-Σ-map-snd $ is-equiv-pos-fun i  f

      step-2 : isEquiv (Wᴰ-in _ _ _ s f  ⊎-map-right (Σ-map-snd $ pos-fun i  f))
      step-2 = isEquiv-∘ (isEquiv-Wᴰ-in _ _ _ s f) step-1

      goal : isEquiv (Wᴰ-in _ _ _ s f  ⊎-map-right (Σ-map-snd $ pos-fun i  f)  equivFun (φ.pos i (s , shape  f)))
      goal = isEquiv-∘ step-2 (equivIsEquiv (φ.pos i _))

    pos :  i  (s : W S (P ))  Q i (shape s)  Pos (μ F) i s
    pos i w .fst = pos-fun i w
    pos i w .snd = is-equiv-pos-fun i w
    {-# INLINE pos #-}

    pos' :  i  (s : W S (P ))  Q i (shape s)  Pos (μ F) i s
    pos' i (sup s f) =
      Q i (φ.shape (s , shape  f))
        ≃⟨ φ.pos i (s , shape  f) 
      (P (just i) s  (Σ[ p  P nothing s ] Q i (shape (f p))))
        ≃⟨ ⊎-right-≃ (Σ-cong-equiv-snd λ p  pos i (f p)) 
      (P (just i) s  (Σ[ p  P nothing s ] Pos (μ F) i (f p)))
        ≃⟨ Wᴰ-in-equiv S (P nothing) (P (just i)) s f 
      Wᴰ S (P nothing) (P (just i)) (sup s f)
        ≃∎

  μ-rec-β' : (G : Container  Ix)
     (α : (F [ G ])  G)
     μ-rec G α  μ-out  [-]-map F (μ-rec G α)  α
  μ-rec-β' G α = ⊸≡-ext (μ-rec.shape-β G α) λ where
    ix (sup s f)  refl

  μ-rec-unique' : (G : Container  Ix)
     (α : (F [ G ])  G)
     isContr (Σ[ α*  μ F  G ] α*  μ-out  [-]-map F α*  α )
  μ-rec-unique' G α .fst .fst = μ-rec G α
  μ-rec-unique' G α .fst .snd = μ-rec-β' G α
  μ-rec-unique' G α .snd (ζ , ζ-β) = ΣPathP (goal , goal-coh) where
    module α = _⊸_ α
    module ζ = _⊸_ ζ

    shape-≡-ext :  w  μ-rec.shape G α w  ζ.shape w
    shape-≡-ext w@(sup s f) = p  q
      module shape-≡-ext where
        p = cong α.shape (ΣPathP (refl′ s , funExt (shape-≡-ext  f)))
        q = sym $ cong _⊸_.shape ζ-β ≡$ w

        filler = compPath-filler p q

    shape-≡ : μ-rec.shape G α  ζ.shape
    shape-≡ = funExt shape-≡-ext

    pos-≡-ext :  ix (w : W S (P ))  PathP  i  Pos G ix (shape-≡ i w)  μ.pos F ix w) (μ-rec.pos G α ix w) (ζ.pos ix w)
    pos-≡-ext ix (sup s f) = compPathP' {B = B} {p = p} {q = q} pᴰ qᴰ
      module pos-≡ where
        B : Shape G  Type _
        B t = Pos G ix t  μ.pos F ix (sup s f)

        open shape-≡-ext s f using (p ; q)

        yᴰ : _  _
        yᴰ = α.pos ix (s , ζ.shape  f) ∙ₑ ⊎-right-≃ (Σ-cong-equiv-snd' (ζ.pos ix  f)) ∙ₑ Wᴰ-in-equiv _ _ _ s f

        pᴰ-fun : PathP  i  Pos G ix (shape-≡-ext.p s f i)  μ.pos F ix (sup s f)) (equivFun $ μ-rec.pos G α ix (sup s f)) (equivFun yᴰ)
        pᴰ-fun i =
          equivFun (α.pos ix (s , shape-≡ i  f))
          ⨟ᴰ ⊎-map-right (Σ-map-snd {B = λ p₁  Pos G ix (shape-≡ i (f p₁))} λ p₁  equivFun (pos-≡-ext ix (f p₁) i))
          ⨟ᴰ Wᴰ-in _ _ _ s f

        pᴰ : PathP  i  B (shape-≡-ext.p s f i)) (μ-rec.pos G α ix (sup s f)) yᴰ
        pᴰ = equivPathP pᴰ-fun

        qᴰ : PathP  i  B (shape-≡-ext.q s f i)) yᴰ (ζ.pos ix (sup s f))
        qᴰ i = ζ-β (~ i) ._⊸_.pos ix (sup s f)

        filler : SquareP  z i  B (compPath-filler p q z i)) pᴰ (compPathP' {B = B} {p = p} {q = q} pᴰ qᴰ) _ qᴰ
        filler = compPathP'-filler {B = B} {p = p} {q = q} pᴰ qᴰ

    pos-≡ : PathP  i   ix w  Pos G ix (shape-≡ i w)  μ.pos F ix w) (μ-rec.pos G α) ζ.pos
    pos-≡ = funExt₂ pos-≡-ext

    goal : μ-rec G α  ζ
    goal = ⊸≡ shape-≡ pos-≡

    goal-coh-shape : Square {A = W S (P nothing)  Shape G}
       i  μ-rec-β' G α i ._⊸_.shape)
       i  ζ-β i ._⊸_.shape)
      shape-≡
      λ i z  α.shape (_⊸_.shape ([-]-map F (goal i)) (_⊸_.shape μ-out z))
    goal-coh-shape i j (sup s f) = shape-≡-ext.filler s f (~ j) i

    goal-coh-pos-ext :  ix w
       SquareP  i j  Pos G ix (goal-coh-shape i j w)  μ.pos F ix w)
         i  μ-rec-β' G α i ._⊸_.pos ix w)
         i  ζ-β i ._⊸_.pos ix w)
         i  pos-≡-ext ix w i)
        λ i  (α.pos ix (_⊸_.shape ([-]-map F (goal i)) (_⊸_.shape μ-out w))) ∙ₑ (_⊸_.pos ([-]-map F (goal i)) ix (_⊸_.shape μ-out w) ∙ₑ _⊸_.pos μ-out ix w)
    goal-coh-pos-ext ix (sup s f) = equivSquareP λ i j  equivFun (pos-≡.filler ix s f (~ j) i)

    goal-coh-pos : SquareP  i j   ix w  Pos G ix (goal-coh-shape i j w)  μ.pos F ix w)
       i  μ-rec-β' G α i ._⊸_.pos)
       i  ζ-β i ._⊸_.pos)
      pos-≡
      λ i ix w  (α.pos ix (_⊸_.shape ([-]-map F (goal i)) (_⊸_.shape μ-out w))) ∙ₑ (_⊸_.pos ([-]-map F (goal i)) ix (_⊸_.shape μ-out w) ∙ₑ _⊸_.pos μ-out ix w)
    goal-coh-pos i j ix w = goal-coh-pos-ext ix w i j

    goal-coh : Square (μ-rec-β' G α) ζ-β goal (cong  f  μ-out  [-]-map F f  α) goal)
    goal-coh i j ._⊸_.shape = goal-coh-shape i j
    goal-coh i j ._⊸_.pos = goal-coh-pos i j

  opaque
    μ-rec-unique : (G : Container  Ix)
       (α : (F [ G ])  G)
       isContr (Σ[ α*  μ F  G ] μ-in  α*  [-]-map F α*  α )
    μ-rec-unique G α = isOfHLevelRespectEquiv 0 (Σ-cong-equiv-snd comm-square-equiv) $ μ-rec-unique' G α
      where module _ (α* : μ F  G) where
        cancel : μ-in  (μ-out  ([-]-map F α*  α))  [-]-map F α*  α
        cancel =
          μ-in  (μ-out  ([-]-map F α*  α))
            ≡⟨ sym (⋆-assoc μ-in μ-out ([-]-map F α*  α)) 
          (μ-in  μ-out)  ([-]-map F α*  α)
            ≡⟨ cong (_⋆ ([-]-map F α*  α)) μ-in-out 
          id (F [ μ F ])  ([-]-map F α*  α)
            ≡⟨ ⋆-id-left _ 
          [-]-map F α*  α
            

        comm-square-equiv : (α*  μ-out  [-]-map F α*  α)  (μ-in  α*  [-]-map F α*  α)
        comm-square-equiv =
          (α*  (μ-out  [-]-map F α*)  α)
            ≃⟨ compPathrEquiv (⋆-assoc μ-out ([-]-map F α*) α) 
          (α*  μ-out  ([-]-map F α*  α))
            ≃⟨ congEquiv (containerEquivPreCompEquiv μ-in-equiv G) 
          (μ-in  α*  μ-in  (μ-out  ([-]-map F α*  α)))
            ≃⟨ compPathrEquiv cancel 
          (μ-in  α*  [-]-map F α*  α)
            ≃∎

    μ-rec-β : (G : Container  Ix)
       (α : (F [ G ])  G)
       μ-in  μ-rec G α  [-]-map F (μ-rec G α)  α
    μ-rec-β G α = μ-rec-unique G α .fst .snd

  isEquivFrom-μ-rec : (G : Container  Ix)
     (φ : (F [ G ])  G)
     isContainerEquiv (μ-rec G φ)
     isContainerEquiv φ
  isEquivFrom-μ-rec G φ is-equiv-μ-rec = is-equiv-φ where
    is-equiv-comp : isContainerEquiv (μ-out  [-]-map F (μ-rec G φ)  φ)
    is-equiv-comp = subst isContainerEquiv (μ-rec-β' G φ) is-equiv-μ-rec

    is-equiv-μ-out⋆F[μ-rec] : isContainerEquiv (μ-out  [-]-map F (μ-rec G φ))
    is-equiv-μ-out⋆F[μ-rec] = isContainerEquivComp
      μ-out
      ([-]-map F (μ-rec G φ))
      (equivIsContainerEquiv μ-out-equiv)
      (isEquiv-[-]-map F (μ-rec G φ) is-equiv-μ-rec)

    is-equiv-φ : isContainerEquiv φ
    is-equiv-φ = isContainerEquivCompRight' (μ-out  [-]-map F (μ-rec G φ)) φ is-equiv-μ-out⋆F[μ-rec] is-equiv-comp

  isEmbedding-μ-rec : (G : Container  Ix)
     (α : (F [ G ])  G)
     isContainerEmbedding α
     isContainerEmbedding (μ-rec G α)
  isEmbedding-μ-rec G α = goal where
    module α = _⊸_ α

    goal : isEmbedding α.shape  isEmbedding (W-rec α.shape)
    goal = isEmbedding-W-rec α.shape