module Polynomial.Composition where

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

open import Polynomial.IR
open import Polynomial.Monad

_IR-/_ : {D E : Set1 }(M : Poly E) → IR D E → IR D (Info M)
idPN IR-/ R = R
con A IR-/ R = con A , id
sigma S F IR-/ R = (S IR-/ R) Polynomial->>=d (λ p → F p IR-/ R)
pi A f IR-/ R = πIR A (λ a → f a IR-/ R)

_/_ : {D E : Set1} -> (F : Poly E)(c : IR D E) -> Poly D
c / R = proj₁ (c IR-/ R)

forget/ : ∀ {D E} (F : Poly E)(c : IR D E) -> Info (F / c) -> Info F
forget/ c R = proj₂ (c IR-/ R)

_○_ : {X Y Z : Set1} -> IR Y Z -> IR X Y -> IR X Z
(S , f) ○ d = (S / d , f ∘ forget/ S d)

mutual

  left0 : {D E : Set1} -> (F : Poly E)(c : IR D E) -> {Z : Fam D} ->
          Node (F / c) Z -> Node F (⟦ c ⟧ Z)
  left0 idPN c x = x
  left0 (con A) c a = a
  left0 (sigma S T) c {Z} (x , y)
    = left0 S c x , subst (λ z → Node (T z) (⟦ c ⟧ Z))
                          (left1 S c x)
                          (left0 (T (forget/ S c (info (S / c) Z x))) c y)
  left0 (pi A T) c f = λ a → left0 (T a) c (f a)

  left1 : {D E : Set1} -> (F : Poly E)(c : IR D E) -> {Z : Fam D} ->
          (x : Node (F / c) Z) ->
          forget/ F c (info (F / c) Z x) ≡ info F (⟦ c ⟧ Z) (left0 F c x)
  left1 idPN c x = refl
  left1 (con A) c a = refl
  left1 (sigma S T) c {Z } (x , y)
     = Σ-≡ (left1 S c x)
           (trans (cong (subst (λ z → Info (T z)) (left1 S c x))
                        (left1 (T (forget/ S c (info (S / c) Z x))) c y))
                  (subst-natural (λ {y} → info (T y) (⟦ c ⟧ Z))
                                 (left1 S c x)))
  left1 (pi A T) c f = ext (λ a → left1 (T a) c (f a))

mutual

  right0 : {D E : Set1} -> (F : Poly E)(c : IR D E) -> {Z : Fam D} ->
           Node F (⟦ c ⟧ Z) -> Node (F / c) Z
  right0 idPN c x = x
  right0 (con A) c a = a
  right0 (sigma S T) c {Z} (x , y)
    = right0 S c x , subst (λ z → Node (T z / c) Z)
                           (right1 S c x)
                           (right0 (T (info S (⟦ c ⟧ Z) x)) c y)
  right0 (pi A T) c f = λ a → right0 (T a) c (f a)

  right1 : {D E : Set1} -> (F : Poly E)(c : IR D E) -> {Z : Fam D} ->
           (x : Node F (⟦ c ⟧ Z)) ->
           info F (⟦ c ⟧ Z) x ≡ forget/ F c (info (F / c) Z (right0 F c x))
  right1 idPN c x = refl
  right1 (con A) c a = refl
  right1 (sigma S T) c {Z} (x , y)
     = Σ-≡ (right1 S c x)
           (trans (cong (subst (λ z → Info (T z)) (right1 S c x))
                        (right1 (T (info S (⟦ c ⟧ Z) x)) c y))
                  (subst-natural (λ {y} → forget/ (T y) c ∘ (info (T y / c) Z))
                                 (right1 S c x)))
  right1 (pi A T) c f = ext (λ a → right1 (T a) c (f a))

------------------------------------------------------

mutual

  leftright0 : {D E : Set1} -> (F : Poly E)(c : IR D E) -> {Z : Fam D} ->
               (x : Node F (⟦ c ⟧ Z)) -> x ≡ left0 F c (right0 F c x)
  leftright0 idPN c x = refl
  leftright0 (con A) c a = refl
  leftright0 (sigma S T) c {Z} (x , y)
     = Σ-≡ (leftright0 S c x)
           (trans
             (subst-cong (leftright0 S c x))
             (trans
               (congSubst (sym (leftright1 S c x)))
               (trans
                 (subst-trans (right1 S c x) (left1 S c (right0 S c x)))
                 (cong (subst (λ v → Node (T v) (⟦ c ⟧ Z))
                              (left1 S c (right0 S c x)))
                       (trans
                         (cong (subst (λ v → Node (T v) (⟦ c ⟧ Z))
                                      (right1 S c x))
                               (leftright0 (T _) c y))
                         (subst-natural (λ {y} → left0 (T y) c)
                                        (right1 S c x))) ))))
  leftright0 (pi A T) c f = ext (λ a → leftright0 (T a) c (f a))

  leftright1 : {D E : Set1} -> (F : Poly E)(c : IR D E) -> {Z : Fam D} ->
               (x : Node F (⟦ c ⟧ Z)) ->
               trans (right1 F c x) (left1 F c (right0 F c x))
                 ≡ cong (info F (⟦ c ⟧ Z)) (leftright0 F c x)
  leftright1 idPN c x = refl
  leftright1 (con A) c a = refl
  leftright1 (sigma S T) c (x , y) = UIP _ _ -- TODO
  leftright1 (pi A T) c {Z} f
     = trans
        (trans-ext (λ a → right1 (T a) c (f a))
                   (λ a → left1 (T a) c (right0 (T a) c (f a))))
        (eq-extensionality (λ a →
         trans
            (ext-β _ a)
            (trans
              (leftright1 (T a) c (f a))
              (trans
                 (cong (cong (info (T a) (⟦ c ⟧ Z))) (sym (ext-β _ a)))
                 ((trans (sym (cong-cong _)) (cong-cong _)))))))

mutual

  rightleft0 : {D E : Set1} -> (F : Poly E)(c : IR D E) -> {Z : Fam D} ->
               (x : Node (F / c) Z) -> x ≡ right0 F c (left0 F c x)
  rightleft0 idPN c x = refl
  rightleft0 (con A) c a = refl
  rightleft0 (sigma S T) c {Z} (x , y)
   = Σ-≡ (rightleft0 S c x)
         (trans
            (subst-cong (rightleft0 S c x))
            (trans
              (congSubst (sym (rightleft1 S c x)))
              (trans
                (subst-trans (left1 S c x) (right1 S c (left0 S c x)))
                (cong (subst (λ z → Node (T z / c) Z)
                             (right1 S c (left0 S c x)))
                      (trans
                        (cong (subst (λ v → Node (T v / c) Z) (left1 S c x))
                              (rightleft0 (T _) c y))
                        (subst-natural (λ {y} → right0 (T y) c)
                                       (left1 S c x)))))))
  rightleft0 (pi A T) c f = ext (λ a → rightleft0 (T a) c (f a))

  rightleft1 : {D E : Set1} -> (F : Poly E)(c : IR D E) -> {Z : Fam D} ->
               (x : Node (F / c) Z) ->
               trans (left1 F c x) (right1 F c (left0 F c x))
                 ≡ cong (forget/ F c ∘ (info (F / c) Z)) (rightleft0 F c x)
  rightleft1 F c x = UIP _ _ -- TODO


○-equiv : {X Y Z : Set1} -> (c : IR Y Z) -> (d : IR X Y) -> {Z : Fam X} ->
          ⟦ c ○ d ⟧ Z ≃ ⟦ c ⟧ (⟦ d ⟧ Z)
○-equiv R@(c , α) d
   = record { left = left0 c d , (cong α ∘ (left1 c d))
            ; right = right0 c d , (cong α ∘ (right1 c d))
            ; leftright = comp-is-id-ext
                            (left0 c d , (cong α ∘ (left1 c d)))
                            (right0 c d , (cong α ∘ (right1 c d)))
                            (leftright0 c d)
                            (λ x → trans
                                      (trans-cong (right1 c d x)
                                                  (left1 c d (right0 c d x)))
                                      (trans
                                        (cong (cong α) (leftright1 c d x))
                                        (sym (cong-cong {f = α}
                                                        (leftright0 c d x)))))
            ; rightleft = comp-is-id-ext
                            (right0 c d , (cong α ∘ (right1 c d)))
                            (left0 c d , (cong α ∘ (left1 c d)))
                            (rightleft0 c d)
                            (λ x → trans
                                      (trans-cong (left1 c d x)
                                                  (right1 c d (left0 c d x)))
                                      (trans
                                        (cong (cong α) (rightleft1 c d x))
                                        (sym (cong-cong {f = α}
                                                        (rightleft0 c d x)))))
            }