module DS.compositionFromPower where


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

open import DS.IR

import DS.Bind as B


------------------------------------------------------------------
postulate
  -- assuming there is a power combinator for Dybjer-Setzer IR codes...
  _⟶_ : {D E : Set1} -> (S : Set) -> DS D E -> DS D (S -> E)
  -- ...with the right decoding...
  ⟶-equiv : ∀ {D E} -> {S : Set}{c : DS D E} -> {Z : Fam D} ->
             ⟦ S ⟶ c ⟧ Z ≃ (S Fam⟶ (⟦ c ⟧ Z))
  -- ...we can define composition of DS codes; see below.

abstract -- it's annoying that these unfold since it makes goals harder to read;
         -- hence we keep them abstract
  ⟶left0 : ∀ {D E} -> {S : Set}{c : DS D E} -> {Z : Fam D} ->
             ⟦ S ⟶ c ⟧₀ Z -> (S -> ⟦ c ⟧₀ Z)
  ⟶left0 = _Fam⇒_.indmor (_≃_.left ⟶-equiv)

  ⟶left1 : ∀ {D E} -> {S : Set}{c : DS D E} -> {Z : Fam D} ->
           (g : ⟦ S ⟶ c ⟧₀ Z) -> ⟦ S ⟶ c ⟧₁ Z g ≡ ⟦ c ⟧₁ Z ∘ ⟶left0 g
  ⟶left1 =  _Fam⇒_.indmor= (_≃_.left ⟶-equiv)

  ⟶right0 : ∀ {D E} -> {S : Set}{c : DS D E} -> {Z : Fam D} ->
            (S -> ⟦ c ⟧₀ Z) -> ⟦ S ⟶ c ⟧₀ Z
  ⟶right0 = _Fam⇒_.indmor (_≃_.right ⟶-equiv)

  ⟶right1 : ∀ {D E} -> {S : Set}{c : DS D E} -> {Z : Fam D} ->
              (g : S -> ⟦ c ⟧₀ Z) -> ⟦ c ⟧₁ Z ∘ g ≡ ⟦ S ⟶ c ⟧₁ Z (⟶right0 g)
  ⟶right1 = _Fam⇒_.indmor= (_≃_.right ⟶-equiv)

  ⟶leftright0 : ∀ {D E} -> {S : Set}{c : DS D E} -> {Z : Fam D} ->
                (g : S -> ⟦ c ⟧₀ Z) -> g ≡ ⟶left0 {c = c} (⟶right0 g)
  ⟶leftright0 g = happly (cong _Fam⇒_.indmor (sym (_≃_.leftright ⟶-equiv))) g

  ⟶rightleft0 : ∀ {D E} -> {S : Set}{c : DS D E} -> {Z : Fam D} ->
                (g : ⟦ S ⟶ c ⟧₀ Z) -> g ≡ ⟶right0 {c = c} (⟶left0 g)
  ⟶rightleft0 g = happly (cong _Fam⇒_.indmor (sym (_≃_.rightleft ⟶-equiv))) g

  ⟶leftright1 : ∀ {D E} -> {S : Set}{c : DS D E} -> {Z : Fam D} ->
                (g : S -> ⟦ c ⟧₀ Z) ->
                trans (⟶right1 {c = c} g) (⟶left1 (⟶right0 g))
                  ≡ cong (λ z → ⟦ c ⟧₁ Z ∘ z) (⟶leftright0 g)
  ⟶leftright1 {S = S} {c} {Z} g
     = trans
         (sym (happly (congd _Fam⇒_.indmor= leftright) g))
         (trans
           (happly (subst-cod leftright) g)
           (trans (subst-path-general leftright) (cong-cong-cong leftright)))
       where leftright = sym (_≃_.leftright ⟶-equiv)

  ⟶rightleft1 : ∀ {D E} -> {S : Set}{c : DS D E} -> {Z : Fam D} ->
                (g : ⟦ S ⟶ c ⟧₀ Z) ->
                trans (⟶left1 {c = c} g) (⟶right1 (⟶left0 g))
                  ≡ cong (⟦ S ⟶ c ⟧₁ Z) (⟶rightleft0 g)
  ⟶rightleft1 {S = S} {c} {Z} g
     = trans
         (sym (happly (congd _Fam⇒_.indmor= rightleft) g))
         (trans
            (happly (subst-cod  rightleft) g)
            (trans (subst-path-general rightleft) (cong-cong-cong rightleft)))
       where rightleft = sym (_≃_.rightleft ⟶-equiv)
------------------------------------------------------------------

_○_ : {C D E : Set1} -> DS D E -> DS C D -> DS C E
ι e ○ d = ι e
σ A f ○ d = σ A (λ a → (f a) ○ d)
δ A F ○ d = (A ⟶ d) DS->>= (λ g → (F g) ○ d)

left0 : {C D E : Set1} -> (c : DS D E) -> (d : DS C D) ->
        {Z : Fam C} -> ⟦ c ○ d ⟧₀ Z -> ⟦ c ⟧₀ (⟦ d ⟧ Z)
left0 (ι e) d _ = _
left0 (σ A f) d (a , x) = a , left0 (f a) d x
left0 (δ A F) d {Z} y
  = let (h , x) = B.left0 (A ⟶ d) _ y
    in (⟶left0 h , subst (λ z → ⟦ F z ⟧₀ (⟦ d ⟧ Z)) (⟶left1 _) (left0 (F _) d x))

left1 : {C D E : Set1} -> (c : DS D E) -> (d : DS C D) ->
        {Z : Fam C} -> (x : ⟦ c ○ d ⟧₀ Z) ->
        ⟦ c ○ d ⟧₁ Z x ≡ ⟦ c ⟧₁ (⟦ d ⟧ Z) (left0 c d x)
left1 (ι e) d x = refl
left1 (σ A f) d (a , x) = left1 (f a) d x
left1 (δ A F) d {Z} y
  = let (h , x) = B.left0 (A ⟶ d) _ y
    in trans
         (B.left1 (A ⟶ d) _ y)
         (trans
           (left1 (F _) d x)
           (happly
             (trans (congd-sym (λ z → ⟦ F z ⟧₁ (⟦ d ⟧ Z)) (⟶left1 h))
                    (subst-dom-sym (⟶left1 h)))
             (left0 (F _) d x)))

right0 : {C D E : Set1} -> (c : DS D E) -> (d : DS C D) ->
        {Z : Fam C} -> ⟦ c ⟧₀ (⟦ d ⟧ Z) -> ⟦ c ○ d ⟧₀ Z
right0 (ι e) d _ = _
right0 (σ A f) d (a , x) = a , right0 (f a) d x
right0 (δ A F) d {Z} (h , x)
  = B.right0 (A ⟶ d) _
             (⟶right0 h , subst (λ z → ⟦ F z ○ d ⟧₀ Z)
                                 (⟶right1 _) (right0 (F _) d x))

right1 : {C D E : Set1} -> (c : DS D E) -> (d : DS C D) ->
         {Z : Fam C} -> (x : ⟦ c ⟧₀ (⟦ d ⟧ Z)) ->
         ⟦ c ⟧₁ (⟦ d ⟧ Z) x ≡ ⟦ c ○ d ⟧₁ Z (right0 c d x)
right1 (ι e) d x = refl
right1 (σ A f) d (a , x) = right1 (f a) d x
right1 (δ A F) d {Z} (h , x)
 = let h' = ⟶right0 {c = d} h
       x' = subst (λ z → ⟦ F z ○ d ⟧₀ Z) (⟶right1 {c = d} _) (right0 (F _) d x)
   in trans
        (trans
          (right1 (F _) d x)
          (happly
            (trans (congd-sym (λ z → ⟦ F z ○ d ⟧₁ Z) (⟶right1 {c = d} h))
                   (subst-dom-sym (⟶right1 h)))
            (right0 (F _) d x)))
        (B.right1 (A ⟶ d) _ (h' , x'))

leftright0 : {C D E : Set1} -> (c : DS D E) -> (d : DS C D) ->
             {Z : Fam C} -> (x : ⟦ c ⟧₀ (⟦ d ⟧ Z)) ->
             x ≡ left0 c d (right0 c d x)
leftright0 (ι e) d _ = refl
leftright0 (σ A f) d (a , x) = cong (_,_ a) (leftright0 (f a) d x)
leftright0 (δ A F) d {Z} (h , x)
 = let h' = ⟶right0 {c = d} h
       x' = subst (λ z → ⟦ F z ○ d ⟧₀ Z) (⟶right1 {c = d} _) (right0 (F _) d x)
       fst-part = cong ⟶left0 (B.leftrightFst0 (A ⟶ d) (λ g → F g ○ d) h' x')
       fst-proof = trans (⟶leftright0 h) fst-part
       lemma : (left0 (F (⟦ d ⟧₁ Z ∘ h)) d (right0 (F (⟦ d ⟧₁ Z ∘ h)) d x))
                 ≡ subst (λ z → ⟦ F z ⟧₀ (⟦ d ⟧ Z))
                     (sym (⟶right1 h))
                     (left0 (F (⟦ A ⟶ d ⟧₁ Z (⟶right0 h))) d
                       (subst (λ z → ⟦ F z ○ d ⟧₀ Z)
                              (⟶right1 h)
                              (right0 (F (⟦ d ⟧₁ Z ∘ h)) d x)))
       lemma = subst-inverse
                 (⟶right1 h)
                 (trans
                   (cong
                     ((subst (λ z → ⟦ F z ⟧₀ (⟦ d ⟧ Z)) (⟶right1 h))
                          ∘ (left0 (F (⟦ d ⟧₁ Z ∘ h)) d))
                     (sym (subst-sym-cancel' (⟶right1 h))))
                   (happly (trans (sym (subst-domcod-simple (⟶right1 h)))
                                  (congd (λ z → left0 (F z) d) (⟶right1 h)))
                           (subst (λ z → ⟦ F z ○ d ⟧₀ Z)
                                  (⟶right1 h)
                                  (right0 (F (⟦ d ⟧₁ Z ∘ h)) d x))))
       path-algebra : trans (sym (⟶right1 {c = d} h))
                            (cong (λ g → ⟦ d ⟧₁ Z ∘ g)
                                  (trans (⟶leftright0 h) fst-part))
                          ≡ trans (⟶left1 (⟶right0 h))
                                  (cong (λ g → ⟦ d ⟧₁ Z ∘ (⟶left0 g))
                                        (B.leftrightFst0
                                           (A ⟶ d)
                                           (λ g → F g ○ d) (⟶right0 h)
                                           (subst (λ z → ⟦ F z ○ d ⟧₀ Z)
                                                  (⟶right1 h)
                                                  (right0 (F (⟦ d ⟧₁ Z ∘ h)) d x))))
       path-algebra = trans
                        (trans
                          (cong (trans (sym (⟶right1 h)))
                                (sym (trans-cong (⟶leftright0 h) fst-part)))
                          (trans
                            (cong (λ z →
                                      trans
                                       (sym (⟶right1 {c = d} h))
                                       (trans z (cong (λ g → ⟦ d ⟧₁ Z ∘ g) fst-part)))
                                  (sym (⟶leftright1 h)))
                            (trans-cancel-adhoc (⟶right1 h)
                                                (⟶left1 (⟶right0 h))
                                                (cong (λ g → ⟦ d ⟧₁ Z ∘ g) fst-part))))
                        (cong (trans (⟶left1 (⟶right0 {c = d} h)))
                              (sym (cong-cong
                                     (B.leftrightFst0
                                           (A ⟶ d)
                                           (λ g → F g ○ d) (⟶right0 h)
                                           (subst (λ z → ⟦ F z ○ d ⟧₀ Z)
                                                  (⟶right1 h)
                                                  (right0 (F (⟦ d ⟧₁ Z ∘ h)) d x))))))
   in Σ-≡ fst-proof
          (trans
            (cong (subst (λ g → ⟦ F (⟦ d ⟧₁ Z ∘ g) ⟧₀ (⟦ d ⟧ Z)) fst-proof)
                  (trans (leftright0 (F _) d x) lemma))
            (trans
              (trans (subst-cong fst-proof)
                     (sym (subst-trans (sym (⟶right1 {c = d} h))
                                       (cong (λ g → ⟦ d ⟧₁ Z ∘ g) fst-proof))))
                (trans
                 (trans
                   (trans
                     (trans (congSubst path-algebra)
                            (subst-trans
                              (⟶left1 (⟶right0 {c = d} h))
                              (cong (λ g → ⟦ d ⟧₁ Z ∘ (⟶left0 g))
                                    (B.leftrightFst0 (A ⟶ d) (λ g → F g ○ d) h' x'))))
                     (sym (subst-cong
                            (B.leftrightFst0 (A ⟶ d) (λ g → F g ○ d) h' x'))))
                   (subst-Σ-≡-proj₁
                     (B.leftrightFst0 (A ⟶ d) (λ g → F g ○ d) h' x')
                     (B.leftrightSnd0 (A ⟶ d) (λ g → F g ○ d) h' x')))
                 (congd
                   (λ z → subst _ (⟶left1 (proj₁ z))
                                   (left0 (F (⟦ A ⟶ d ⟧₁ Z (proj₁ z))) d (proj₂ z)))
                   (B.leftright0 (A ⟶ d) (λ g → F g ○ d) (h' , x'))))))

rightleft0 : {C D E : Set1} -> (c : DS D E) -> (d : DS C D) ->
             {Z : Fam C} -> (x : ⟦ c ○ d ⟧₀ Z) ->
             x ≡ right0 c d (left0 c d x)
rightleft0 (ι e) d _ = refl
rightleft0 (σ A f) d (a , x) = cong (_,_ a) (rightleft0 (f a) d x)
rightleft0 (δ A F) d {Z} y
  = let (h , x) = B.left0 (A ⟶ d) _ y
    in trans
         (B.rightleft0 (A ⟶ d) (λ g → F g ○ d) y)
         (cong (B.right0 (A ⟶ d) (λ g → F g ○ d))
               (Σ-≡ (⟶rightleft0 h)
                    (trans
                      (subst-cong (⟶rightleft0 h))
                      (trans
                        (congSubst (sym (⟶rightleft1 h)))
                        (trans
                          (subst-trans (⟶left1 h) (⟶right1 (⟶left0 h)))
                          (cong
                            (subst (λ z → ⟦ F z ○ d ⟧₀ Z) (⟶right1 (⟶left0 h)))
                            (trans
                              (cong (subst (λ z → ⟦ F z ○ d ⟧₀ Z) (⟶left1 h))
                                    (trans
                                      (rightleft0 (F _) d x)
                                      (cong (right0 (F _) d)
                                            (sym (subst-sym-cancel' (⟶left1 h))))))
                              (happly (trans
                                        (sym (subst-domcod-simple (⟶left1 h)))
                                        (congd (λ z → right0 (F z) d)
                                               (⟶left1 {c = d} h)))
                                      (subst (λ z → ⟦ F z ⟧₀ (⟦ d ⟧ Z))
                                             (⟶left1 h)
                                             (left0 (F (⟦ A ⟶ d ⟧₁ Z h)) d x))))))))))

leftright1 : {C D E : Set1} -> (c : DS D E) -> (d : DS C D) ->
             {Z : Fam C} -> (x : ⟦ c ⟧₀ (⟦ d ⟧ Z)) ->
             trans (right1 c d x) (left1 c d (right0 c d x))
                ≡ cong (⟦ c ⟧₁ (⟦ d ⟧ Z)) (leftright0 c d x)
leftright1 (ι e) d x = refl
leftright1 (σ A f) d (a , x) = trans (leftright1 (f a) d x)
                                     (cong-cong (leftright0 (f a) d x))
leftright1 (δ A F) d x = UIP _ _ -- TODO: this is a big goal

rightleft1 : {C D E : Set1} -> (c : DS D E) -> (d : DS C D) ->
             {Z : Fam C} -> (x : ⟦ c ○ d ⟧₀ Z) ->
             trans (left1 c d x) (right1 c d (left0 c d x))
                ≡ cong (⟦ c ○ d ⟧₁ Z) (rightleft0 c d x)
rightleft1 (ι e) d x = refl
rightleft1 (σ A f) d {Z} (a , x) = trans (rightleft1 (f a) d x)
                                         (cong-cong (rightleft0 (f a) d x))
rightleft1 (δ A F) d x = UIP _ _ -- TODO: this is a big goal

-- Combining it all into an equivalence

○-Equiv : {C D E : Set1} -> (c : DS D E) -> (d : DS C D) ->
          {Z : Fam C} -> ⟦ c ○ d ⟧ Z ≃ ⟦ c ⟧ (⟦ d ⟧ Z)
○-Equiv c d = record { left = (left0 c d , left1 c d)
                     ; right = (right0 c d , right1 c d)
                     ; leftright = comp-is-id-ext (left0 c d , left1 c d)
                                                  (right0 c d , right1 c d)
                                                  (leftright0 c d)
                                                  (leftright1 c d)
                     ; rightleft = comp-is-id-ext (right0 c d , right1 c d)
                                                  (left0 c d , left1 c d)
                                                  (rightleft0 c d)
                                                  (rightleft1 c d)
                     }