module Polynomial.Container.Composition where

open import Prelude.Basic
open import Prelude.Equality
open import Prelude.Fam

open import Polynomial.Container.IR

_/_ : {D E : Set₁}(M : IRCont E) -> IR D E -> IR D (IRPos M)
(eta (S , P)) / F  = (con S , proj₁) IRDbind (λ s -> πIR (P s) (const F))
(sigma R f) / F = (R / F) IRDbind (λ p -> (f p) / F)
(π A f) / F  = πIR A (λ a → f a / F)

_/ix_ : {D E : Set₁}(M : IRCont E) -> IR D E -> IRCont D
R /ix F = proj₁ (R / F)

_/dx_ : {D E : Set₁}(M : IRCont E)(F : IR D E) -> IRPos (M /ix F) -> IRPos M
R /dx F = proj₂ (R / F)

_○_ : {C D E : Set1} -> IR D E -> IR C D -> IR C E
(Q , α) ○ G = IR-map α (Q / G)

{- Proof that ○ really is composition -}

-- First, maps back and forth
mutual
  ix-right : {C D : Set1} -> (F : IRCont D) -> (G : IR C D) -> (X : Fam C) ->
          ix (F /ix G) X -> ix F (⟦ G ⟧ X)
  ix-right (eta (S , P)) G X ((s , _) , g) = s , g
  ix-right (sigma R f) G X (p , q)
    = ix-right R G X p , subst (λ z → ix (f z) (⟦ G ⟧ X))
                               (dx-right R G X p)
                               (ix-right (f ((R /dx G) (dx (R /ix G) X p))) G X q)
  ix-right (π A f) G X g a = ix-right (f a) G X (g a)

  dx-right : {C D : Set1} -> (F : IRCont D) -> (G : IR C D) ->
           (X : Fam C) -> (y : ix (F /ix G) X) ->
           (F /dx G) (dx (F /ix G) X y) ≡ dx F (⟦ G ⟧ X) ((ix-right F G X) y)
  dx-right (eta (S , P)) G X ((s , _) , g) = refl
  dx-right (sigma R f) G X (p , q) rewrite dx-right R G X p
    = cong₂d _,_ refl (dx-right (f (dx R (⟦ G ⟧ X) (ix-right R G X p))) G X q)
  dx-right (π A f) G X g = ext (λ a → dx-right (f a) G X (g a))

mutual
  ix-left : {C D : Set1} -> (F : IRCont D) -> (G : IR C D) -> (X : Fam C) ->
          ix F (⟦ G ⟧ X) -> ix (F /ix G) X
  ix-left (eta (S , P)) G X (s , g) = (s , (λ ())) , g
  ix-left (sigma R f) G X (p , q)
    = (ix-left R G X p) , subst (λ z → ix (f z /ix G) X)
                                (dx-left R G X p)
                                (ix-left (f (dx R (⟦ G ⟧ X) p)) G X q)
  ix-left (π A f) G X g a = ix-left (f a) G X (g a)

  dx-left : {C D : Set1} -> (F : IRCont D) -> (G : IR C D) ->
           (X : Fam C) -> (y : ix F (⟦ G ⟧ X)) ->
           dx F (⟦ G ⟧ X) y ≡ (F /dx G) (dx (F /ix G) X (ix-left F G X y))
  dx-left (eta (S , P)) G X (s , g) = refl
  dx-left (sigma R f) G X (p , q) rewrite dx-left R G X p
    = cong₂d _,_ refl
                 (dx-left (f ((R /dx G) (dx (R /ix G) X (ix-left R G X p)))) G X q)
  dx-left (π A f) G X g = ext (λ a → dx-left (f a) G X (g a))

-- Maps back and forth --------------------
right : {C D E : Set1} -> (F : IR D E) -> (G : IR C D) -> (X : Fam C) ->
       ⟦ F ○ G ⟧ X Fam⇒ ⟦ F ⟧ (⟦ G ⟧ X)
right (F , α) G X = ix-right F G X , (λ y → cong α (dx-right F G X y))

left : {C D E : Set1} -> (F : IR D E) -> (G : IR C D) -> (X : Fam C) ->
       ⟦ F ⟧ (⟦ G ⟧ X) Fam⇒ ⟦ F ○ G ⟧ X
left (F , α) G X = ix-left F G X , (λ y → cong α (dx-left F G X y))
-------------------------------------------

-- The roundtrips are identities
mutual
  leftright-ix : {C D : Set1} ->
                 (F : IRCont D) -> (G : IR C D) -> (X : Fam C) ->
                 (x : ix (F /ix G) X) -> x ≡ ix-left F G X (ix-right F G X x)
  leftright-ix (eta (S , P)) G X ((s , _) , g)
    = cong₃-half-d (λ x y z → ((x , y) , z)) refl (ext (λ ())) refl
  leftright-ix (sigma R f) G X (p , q) = cong₂d _,_ (leftright-ix R G X p) calc
    where
     calc = begin
           subst (λ z → ix (f ((R /dx G) (dx (R /ix G) X z)) /ix G) X)
                 (leftright-ix R G X p) q
           ≡⟨ subst-cong (leftright-ix R G X p) ⟩
           subst (λ z → ix (f z /ix G) X)
                 (cong ((R /dx G) ∘ (dx (R /ix G) X)) (leftright-ix R G X p))
                 q
           ≡⟨ cong (λ w → subst (λ z → ix (f z /ix G) X) w q)
                  (sym (leftright-dx R G X p)) ⟩
           subst (λ z → ix (f z /ix G) X)
                 (trans (dx-right R G X p) (dx-left R G X (ix-right R G X p)))
                 q
           ≡⟨ subst-trans (dx-right R G X p) (dx-left R G X (ix-right R G X p)) ⟩
           subst (λ z → ix (f z /ix G) X)
                 (dx-left R G X (ix-right R G X p))
                 (subst (λ z → ix (f z /ix G) X)
                        (dx-right R G X p)
                        q)
           ≡⟨ cong ((subst (λ z → ix (f z /ix G) X)
                          (dx-left R G X (ix-right R G X p))) ∘
                    (subst (λ z → ix (f z /ix G) X)
                           (dx-right R G X p)))
                    (leftright-ix (f ((R /dx G) (dx (R /ix G) X p))) G X q) ⟩
           subst (λ z → ix (f z /ix G) X)
                 (dx-left R G X (ix-right R G X p))
                 (subst (λ z → ix (f z /ix G) X)
                        (dx-right R G X p)
                        (ix-left (f ((R /dx G) (dx (R /ix G) X p)))
                                  G X
                                  (ix-right (f ((R /dx G) (dx (R /ix G) X p)))
                                            G X q)))
           ≡⟨ cong (subst (λ z → ix (f z /ix G) X)
                         (dx-left R G X (ix-right R G X p)))
                   (subst-natural (λ {y} → ix-left (f y) G X)
                                  (dx-right R G X p)) ⟩
           subst (λ z → ix (f z /ix G) X)
                 (dx-left R G X (ix-right R G X p))
                 (ix-left (f (dx R (⟦ G ⟧ X) (ix-right R G X p)))
                           G X
                           (subst (λ z → ix (f z) (⟦ G ⟧ X))
                                  (dx-right R G X p)
                                  (ix-right (f ((R /dx G) (dx (R /ix G) X p)))
                                            G X q)))
           ∎ where open ≡-Reasoning
  leftright-ix (π A f) G X g = ext (λ a → leftright-ix (f a) G X (g a))

  leftright-dx : {C D : Set1} ->
                 (F : IRCont D) -> (G : IR C D) -> (X : Fam C) ->
                 (x : ix (F /ix G) X) ->
                 trans (dx-right F G X x) (dx-left F G X (ix-right F G X x))
                    ≡ cong ((F /dx G) ∘ (dx (F /ix G) X)) (leftright-ix F G X x)
  leftright-dx F G X x = UIP _ _ -- for now

mutual
  rightleft-ix : {C D : Set1} ->
                 (F : IRCont D) -> (G : IR C D) -> (X : Fam C) ->
                 (x : _) -> x ≡ ix-right F G X (ix-left F G X x)
  rightleft-ix (eta (S , P)) G X (s , g) = refl
  rightleft-ix (sigma R f) G X (p , q) = cong₂d _,_ (rightleft-ix R G X p) calc
    where
     calc = begin
           subst (λ z → ix (f (dx R (⟦ G ⟧ X) z)) (⟦ G ⟧ X))
                 (rightleft-ix R G X p)
                 q
           ≡⟨ subst-cong {f = dx R (⟦ G ⟧ X)} (rightleft-ix R G X p) {u = q} ⟩
           subst (λ z → ix (f z) (⟦ G ⟧ X))
                 (cong (dx R (⟦ G ⟧ X)) (rightleft-ix R G X p))
                        q
           ≡⟨ cong (λ w → subst (λ z → ix (f z) (⟦ G ⟧ X)) w q)
                  (sym (rightleft-dx R G X p)) ⟩
           subst (λ z → ix (f z) (⟦ G ⟧ X))
                 (trans (dx-left R G X p) (dx-right R G X (ix-left R G X p)))
                 q

           ≡⟨ subst-trans (dx-left R G X p) (dx-right R G X (ix-left R G X p)) ⟩
           subst (λ z → ix (f z) (⟦ G ⟧ X))
                 (dx-right R G X (ix-left R G X p))
                 (subst (λ z → ix (f z) (⟦ G ⟧ X))
                        (dx-left R G X p)
                        q)
           ≡⟨ cong ((subst (λ z → ix (f z) (⟦ G ⟧ X))
                          (dx-right R G X (ix-left R G X p))) ∘
                    (subst (λ z → ix (f z) (⟦ G ⟧ X))
                           (dx-left R G X p)))
                    (rightleft-ix (f (dx R (⟦ G ⟧ X) p)) G X q) ⟩
           subst (λ z → ix (f z) (⟦ G ⟧ X))
                 (dx-right R G X (ix-left R G X p))
                 (subst (λ z → ix (f z) (⟦ G ⟧ X))
                        (dx-left R G X p)
                        (ix-right (f (dx R (⟦ G ⟧ X) p))
                                  G X
                                  (ix-left (f (dx R (⟦ G ⟧ X) p)) G X q)))
           ≡⟨ cong (subst (λ z → ix (f z) (⟦ G ⟧ X))
                         (dx-right R G X (ix-left R G X p)))
                   (subst-natural (λ {y} → ix-right (f y) G X)
                                  (dx-left R G X p) ) ⟩
           subst (λ z → ix (f z) (⟦ G ⟧ X))
                 (dx-right R G X (ix-left R G X p))
                 (ix-right (f ((R /dx G) (dx (R /ix G) X (ix-left R G X p))))
                           G X
                           (subst (λ z → ix (f z /ix G) X)
                                  (dx-left R G X p)
                                  (ix-left (f (dx R (⟦ G ⟧ X) p)) G X q)))
           ∎ where open ≡-Reasoning
  rightleft-ix (π A f) G X g = ext (λ a → rightleft-ix (f a) G X (g a))


  rightleft-dx : {C D : Set1} ->
                 (F : IRCont D) -> (G : IR C D) -> (X : Fam C) ->
                 (x : _) ->
                 trans (dx-left F G X x) (dx-right F G X (ix-left F G X x))
                    ≡ cong (dx F (⟦ G ⟧ X)) (rightleft-ix F G X x)
  rightleft-dx (eta (S , P)) G X _ = refl
  rightleft-dx (sigma R f) G X (p , q) = UIP _ _ -- TODO: for now
  rightleft-dx (π A f) G X g =
    begin
      trans (ext (λ a → dx-left (f a) G X (g a)))
            (ext (λ a → dx-right (f a) G X (ix-left (f a) G X (g a))))
    ≡⟨ trans-ext _ _ ⟩
      ext (λ x → trans (dx-left (f x) G X (g x))
                        (dx-right (f x) G X (ix-left (f x) G X (g x))))
    ≡⟨ eq-extensionality pp ⟩
      cong (λ g₁ a₁ → dx (f a₁) (⟦ G ⟧ X) (g₁ a₁))
           (ext (λ a₁ → rightleft-ix (f a₁) G X (g a₁)))
    ∎ where
        open ≡-Reasoning
        pp : (a : _) -> _ ≡ _
        pp a = begin
                happly
                  (ext (λ v → trans (dx-left (f v) G X (g v))
                                     (dx-right (f v) G X
                                               (ix-left (f v) G X (g v))))) a
               ≡⟨ ext-β _ a ⟩
                trans (dx-left (f a) G X (g a))
                      (dx-right (f a) G X (ix-left (f a) G X (g a)))
               ≡⟨ rightleft-dx (f a) G X (g a) ⟩
                cong (dx (f a) (⟦ G ⟧ X)) (rightleft-ix (f a) G X (g a))
               ≡⟨ cong (cong (dx (f a) (⟦ G ⟧ X))) (sym (ext-β _ a)) ⟩
                cong (dx (f a) (⟦ G ⟧ X))
                     (happly (ext (λ a₁ → rightleft-ix (f a₁) G X (g a₁))) a)
               ≡⟨ sym (cong-cong _) ⟩
                cong ((λ z → z a) ∘ (λ g₁ a₁ → dx (f a₁) (⟦ G ⟧ X) (g₁ a₁)))
                     (ext (λ a₁ → rightleft-ix (f a₁) G X (g a₁)))
               ≡⟨ cong-cong _ ⟩
                happly (cong (λ g₁ a₁ → dx (f a₁) (⟦ G ⟧ X) (g₁ a₁))
                       (ext (λ a₁ → rightleft-ix (f a₁) G X (g a₁)))) a
               ∎


--Summing up: roundtrips are identities --------------------
rightleft : {C D E : Set1} -> (F : IR D E) -> (G : IR C D) -> (X : Fam C) ->
            (right F G X) Fam∘ (left F G X) ≡ FamId (⟦ F ⟧ (⟦ G ⟧ X))
rightleft (F , α) G X = comp-is-id (right (F , α) G X) (left (F , α) G X)
                                   (ext (rightleft-ix F G X)) small-calc
  where small-calc : (y : _) -> _
        small-calc y = begin
                         trans (cong α (dx-left F G X y))
                               (cong α (dx-right F G X (ix-left F G X y)))
                       ≡⟨ trans-cong (dx-left F G X y)
                                    (dx-right F G X (ix-left F G X y)) ⟩
                         cong α (trans (dx-left F G X y)
                                       (dx-right F G X (ix-left F G X y)))
                       ≡⟨ cong (cong α) (rightleft-dx F G X y) ⟩
                         cong α (cong (dx F (⟦ G ⟧ X))
                                      (rightleft-ix F G X y))
                       ≡⟨ sym (cong-cong (rightleft-ix F G X y)) ⟩
                         cong (α ∘ (dx F (⟦ G ⟧ X)))
                              (rightleft-ix F G X y)
                       ≡⟨ cong (cong (α ∘ (dx F (⟦ G ⟧ X))))
                              (sym (ext-β (rightleft-ix F G X) y)) ⟩
                         cong (α ∘ (dx F (⟦ G ⟧ X)))
                              (happly (ext (rightleft-ix F G X)) y)
                       ∎ where open ≡-Reasoning

leftright : {C D E : Set1} -> (F : IR D E) -> (G : IR C D) -> (X : Fam C) ->
            (left F G X) Fam∘ (right F G X)  ≡ FamId (⟦ F ○ G ⟧ X)
leftright (F , α) G X = comp-is-id (left (F , α) G X) (right (F , α) G X)
                                   (ext (leftright-ix F G X)) small-calc
  where small-calc : (y : _) -> _
        small-calc y = begin
                         trans (cong α (dx-right F G X y))
                               (cong α (dx-left F G X (ix-right F G X y)))
                       ≡⟨ trans-cong (dx-right F G X y)
                                    (dx-left F G X (ix-right F G X y)) ⟩
                         cong α (trans (dx-right F G X y)
                                       (dx-left F G X (ix-right F G X y)))
                       ≡⟨ cong (cong α) (leftright-dx F G X y) ⟩
                         cong α (cong ((F /dx G) ∘ (dx (F /ix G) X))
                                      (leftright-ix F G X y))
                       ≡⟨ sym (cong-cong (leftright-ix F G X y)) ⟩
                         cong (α ∘ (F /dx G) ∘ (dx (F /ix G) X))
                              (leftright-ix F G X y)
                       ≡⟨ cong (cong (α ∘ (F /dx G) ∘ (dx (F /ix G) X)))
                              (sym (ext-β (leftright-ix F G X) y)) ⟩
                         cong (α ∘ (F /dx G) ∘ (dx (F /ix G) X))
                              (happly (ext (leftright-ix F G X)) y)
                       ∎ where open ≡-Reasoning
------------------------------------------------------