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