module Uniform.Composition where

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

open import Uniform.IR
import Uniform.AppLift as AL

mutual
  _○Uni_ : {C D : Set1} -> Uni D -> UF C D -> Uni C
  ι ○Uni R = ι
  (σ c A) ○Uni R = σ (c ○Uni R) (A  [ c ○Uni R ]-Info)
  (δ c A) ○Uni R@(d , β) = (c ○Uni R) AL.++[ (A  [ c ○Uni R ]-Info)  d ]

  [_○Uni_]-Info : {C D : Set1} -> (c : Uni D) -> (R : UF C D) ->
                 Info (c ○Uni R) -> Info c
  [ ι ○Uni R ]-Info x = x
  [ σ c A ○Uni R ]-Info (x , y) = ([ c ○Uni R ]-Info x , y)
  [ δ c A ○Uni R@(d , β) ]-Info x = [ c ○Uni R ]-Info (AL.fst d x) , (β  (AL.snd d x))

_○_ : {C D E : Set1} -> UF D E -> UF C D -> UF C E
(c , α)  R = (c ○Uni R , α  [ c ○Uni R ]-Info)

mutual

  right0 : {C D : Set1} -> (c : Uni D) -> (Q : UF C D) ->
          {Z : Fam C} ->  c ○Uni Q ⟧Uni Z ->  c ⟧Uni ( Q  Z)
  right0 ι Q x = x
  right0 (σ c A) Q (x , a) = (right0 c Q x ,  subst A (right1 c Q x) a)
  right0 (δ c A) Q@(d , β) y
    = let x = AL.rightFst0 d y
          g = AL.rightSnd0 d y
      in right0 c Q x ,  a  g (subst A (sym (right1 c Q x)) a))

  right1 : {C D : Set1} -> (c : Uni D) -> (Q : UF C D) ->
          {Z : Fam C} -> (x :  c ○Uni Q ⟧Uni Z) ->
          [_○Uni_]-Info c Q ( c ○Uni Q ⟧Info Z x)
                ( c ⟧Info ( Q  Z)) (right0 c Q x)
  right1 ι Q x = refl
  right1 (σ c A) Q (x , g) = Σ-≡ (right1 c Q x) refl
  right1 (δ c A) Q@(d , β) {Z} y
    = let x = AL.rightFst0 d y
          g = AL.rightSnd0 d y
          x-lemma = cong ([ c ○Uni Q ]-Info) (AL.rightFst1 d y)
          g-lemma = AL.rightSnd1 d y
      in
        Σ-≡ (trans x-lemma (right1 c Q x))
            (trans (subst-trans x-lemma (right1 c (d , β) x))
              (trans (subst-dom (right1 c (d , β) x)) (ext  a 
                trans (happly (subst-dom x-lemma) _)
                 (cong β (trans
                   (cong (AL.snd d ( _ AL.++[ A  [ c ○Uni Q ]-Info  d ] ⟧Info Z y))
                         (sym (subst-sym-cong (AL.rightFst1 d y))))
                    (g-lemma (subst A (sym (right1 c Q x)) a))))))))


----------------------------------------
right : {C D E : Set1} -> (R : UF D E) -> (Q : UF C D) ->
        {Z : Fam C} ->  R  Q  Z Fam⇒  R  ( Q  Z)
right (c , α) Q = (right0 c Q , (cong α)  (right1 c Q))
----------------------------------------

mutual

  left0 : {C D : Set1} -> (c : Uni D) -> (Q : UF C D) ->
          {Z : Fam C} ->  c ⟧Uni ( Q  Z) ->  c ○Uni Q ⟧Uni Z
  left0 ι Q x = x
  left0 (σ c A) Q (x , a) = (left0 c Q x ,  subst A (left1 c Q x) a)
  left0 (δ c A) Q@(d , β) (x , g)
    = let x' = left0 c Q x
          g' = g  (subst A (sym (left1 c Q x)))
      in
          AL.left0 d (x' , g')

  left1 : {C D : Set1} -> (c : Uni D) -> (Q : UF C D) ->
          {Z : Fam C} -> (x :  c ⟧Uni ( Q  Z)) ->
           c ⟧Info ( Q  Z) x
               [ c ○Uni Q ]-Info ( c ○Uni Q ⟧Info Z (left0 c Q x))
  left1 ι Q x = refl
  left1 (σ c A) Q (x , g) = Σ-≡ (left1 c Q x) refl
  left1 {C} {D} (δ c A) Q@(d , β) {Z} (x , g)
    = let x' = left0 c Q x
          g' = g  (subst A (sym (left1 c Q x)))
          lemma-fst = AL.leftFst1 d x' g'
      in
        Σ-≡ (trans (left1 c Q x) (cong ([ c ○Uni Q ]-Info) lemma-fst))
            -- surely this can be simplified?
            (trans (subst-dom (trans (left1 c Q x)
                                     (cong [ c ○Uni Q ]-Info lemma-fst)))
                   (ext  e  cong β
                     (trans
                       (cong ( d ⟧Info Z  g)
                             (trans
                                (subst-sym-trans (left1 c Q x) _)
                                (cong (subst A (sym (left1 c Q x)))
                                      (trans
                                        (congSubst (sym (sym-cong lemma-fst)))
                                        (sym (subst-cong (sym lemma-fst)))))))
                       (AL.leftSnd1 d x' g' e)))))

----------------------------------------
left : {C D E : Set1} -> (R : UF D E) -> (Q : UF C D) ->
       {Z : Fam C} ->   R  ( Q  Z) Fam⇒  R  Q  Z
left (c , α) Q = left0 c Q ,  x  cong α (left1 c Q x))
----------------------------------------

mutual
  leftright0 : {C D : Set1} -> (c : Uni D) -> (Q : UF C D) ->
               {Z : Fam C} -> (x :  c ○Uni Q ⟧Uni Z) ->
               x  left0 c Q (right0 c Q x)
  leftright0 ι Q x = refl
  leftright0 (σ c A) Q (x , a)
   = Σ-≡ (leftright0 c Q x)
         (trans (subst-cong {B = A} (leftright0 c Q x))
                (trans (congSubst (sym (leftright1 c Q x)))
                       (subst-trans (right1 c Q x) (left1 c Q (right0 c Q x)))))
  leftright0 (δ c A) Q@(d , β) y
   = trans (AL.leftright0 d y)
           (cong (AL.left0 d)
                 (Σ-≡ (leftright0 c Q (AL.rightFst0 d y))
                      (trans
                        (subst-dom
                          (leftright0 c Q (AL.rightFst0 d y)))
                          (ext  a 
                                cong (AL.rightSnd0 d y)
                                     (trans
                                       (subst-cong {B = A}
                                         (sym (leftright0 c Q (AL.rightFst0 d y))))
                                       (trans
                                         (congSubst
                                           (sym-cong (leftright0 c Q
                                                                 (AL.rightFst0 d y))))
                                         (trans
                                           (cong  z  subst A (sym z) a)
                                                 (sym (leftright1 c Q
                                                                  (AL.rightFst0 d y))))
                                           (subst-sym-trans
                                             (right1 c Q (AL.rightFst0 d y))
                                             (left1 c Q
                                                    (right0 c Q
                                                            (AL.rightFst0 d y))))))))))))

  leftright1 : {C D : Set1} -> (c : Uni D) -> (Q : UF C D) ->
               {Z : Fam C} -> (x :  c ○Uni Q ⟧Uni Z) ->
               trans (right1 c Q x) (left1 c Q (right0 c Q x))
                     cong (([ c ○Uni Q ]-Info)  ( c ○Uni Q ⟧Info Z)) (leftright0 c Q x)
  leftright1 c Q x = UIP _ _ -- TODO: for now

mutual
  rightleft0 : {C D : Set1} -> (c : Uni D) -> (Q : UF C D) ->
               {Z : Fam C} -> (x :  c ⟧Uni ( Q  Z)) ->
               x  right0 c Q (left0 c Q x)
  rightleft0 ι Q x = refl
  rightleft0 (σ c A) Q@(d , β) (x , a)
     = Σ-≡ (rightleft0 c Q x)
           (trans (subst-cong {B = A} (rightleft0 c Q x))
                  (trans (congSubst (rightleft1 c Q x))
                         (subst-trans (left1 c Q x) (right1 c Q (left0 c Q x)))))
  rightleft0 {C} {D} (δ c A) Q@(d , β) {Z} y@(x , g)
    = let x' = left0 c Q x
          g' = g  (subst A (sym (left1 c Q x)))
          AL-fst = (cong (right0 c Q) (AL.rightleftFst0 d x' g'))
          proof-fst = (trans (rightleft0 c Q x) AL-fst)
          another-part = sym (right1 c Q
                                      (AL.rightFst0 {c = c ○Uni Q} d
                                                    (AL.left0 d (x' , g'))))
      in
         Σ-≡ proof-fst
             (trans
               (subst-dom proof-fst)
               (ext  a 
                       trans
                        (cong g
                              (trans
                                (congSubst
                                  (trans-sym {x = x}
                                             {y = right0 c Q x'}
                                             (rightleft0 c Q x) AL-fst))
                                (trans
                                  (subst-trans (sym AL-fst) (sym (rightleft0 c Q x)))
                                  (trans
                                    (subst-cong {B = A} (sym (rightleft0 c Q x)))
                                      (trans
                                        (congSubst (rightleft1-sym c Q x))
                                        (trans
                                          (subst-trans (sym (right1 c Q x'))
                                                       (sym (left1 c Q x)))
                                        (cong
                                          (subst A (sym (left1 c Q x)))
                                          (trans
                                            (cong
                                              (subst A
                                                     (sym (right1 c Q x')))
                                              (subst-cong {B = A} (sym AL-fst)))
                                            (trans
                                              (sym
                                               (subst-trans
                                                 (cong ( c ⟧Info ( Q  Z))
                                                       (sym AL-fst))
                                                 (sym (right1 c Q x'))))
                                              (trans
                                                (trans
                                                  (congSubst {B = A}
                                                    (trans
                                                      (cong
                                                        (flip trans
                                                              (sym (right1 c Q x')))
                                                        (sym (trans
                                                         (cong-cong
                                                           (sym
                                                             (AL.rightleftFst0 d x' g')))
                                                         (cong
                                                          (cong ( c ⟧Info ( Q  Z)))
                                                          (sym-cong
                                                            (AL.rightleftFst0 d
                                                                              x' g'))))))
                                                      (trans
                                                        (cong-natural
                                                         (sym  right1 c Q)
                                                         (sym
                                                           (AL.rightleftFst0 d x' g')))
                                                       (cong
                                                         (trans another-part)
                                                         (cong-cong
                                                           (sym
                                                             (AL.rightleftFst0 d
                                                                               x'
                                                                               g'))))))
                                                     {a})
                                                  (subst-trans {B = A}
                                                    another-part
                                                    (cong
                                                      [ c ○Uni Q ]-Info
                                                      (cong
                                                        ( c ○Uni Q ⟧Info Z)
                                                        (sym
                                                          (AL.rightleftFst0 d x' g'))))
                                                    {a}))
                                                (sym
                                                  (subst-cong {B = A}
                                                              {f = [ c ○Uni d , β ]-Info}
                                                    (cong
                                                      ( c ○Uni Q ⟧Info Z)
                                                      (sym (AL.rightleftFst0 d x' g')))
                                                    {subst A another-part a}))))))))))))
                                  (AL.rightleftSnd0 d x' g' (subst A another-part a)))))


  rightleft1 : {C D : Set1} -> (c : Uni D) -> (Q : UF C D) ->
               {Z : Fam C} -> (x :  c ⟧Uni ( Q  Z)) ->
               cong ( c ⟧Info ( Q  Z)) (rightleft0 c Q x)
                     trans (left1 c Q x) (right1 c Q (left0 c Q x))
  rightleft1 c Q x = UIP _ _ -- TODO: for now

  rightleft1-sym : {C D : Set1} -> (c : Uni D) -> (Q : UF C D) ->
                   {Z : Fam C} -> (x :  c ⟧Uni ( Q  Z)) ->
                   cong ( c ⟧Info ( Q  Z)) (sym (rightleft0 c Q x))
                         trans (sym (right1 c Q (left0 c Q x))) (sym (left1 c Q x))
  rightleft1-sym c Q {Z} x = trans (sym-cong (rightleft0 c Q x))
                                   (trans (cong sym (rightleft1 c Q x))
                                          (trans-sym {x =  c ⟧Info ( Q  Z) x}
                                                     {y = [ c ○Uni Q ]-Info
                                                              ( c ○Uni Q ⟧Info Z
                                                                   (left0 c Q x))}
                                                     (left1 c Q x)
                                                     (right1 c Q (left0 c Q x))))

----------------------------------------
rightleft : {C D E : Set1} -> (R : UF D E) -> (Q : UF C D) ->
            {Z : Fam C} ->
            (right R Q {Z} Fam∘ left R Q)  FamId _
rightleft R@(c , α) Q {Z}
  = comp-is-id-ext (right R Q {Z}) (left R Q)
                   (rightleft0 c Q {Z})
                    x 
                     trans
                       (trans-cong (left1 c Q x) (right1 c Q (left0 c Q x)))
                       (trans
                         (cong (cong α) (sym (rightleft1 c Q x)))
                         (sym (cong-cong (rightleft0 c Q x)))))

leftright : {C D E : Set1} -> (R : UF D E) -> (Q : UF C D) ->
            {Z : Fam C} ->
            (left R Q {Z} Fam∘ right R Q)  FamId _
leftright R@(c , α) Q {Z}
  = comp-is-id-ext (left R Q {Z}) (right R Q)
                   (leftright0 c Q {Z})
                    x 
                     trans
                       (trans-cong (right1 c Q x) (left1 c Q (right0 c Q x)))
                       (trans
                         (cong (cong α) (leftright1 c Q x))
                         (sym (cong-cong (leftright0 c Q x)))))

○-equiv : {C D E : Set1} -> (R : UF D E) -> (Q : UF C D) ->
          {Z : Fam C} ->  R  ( Q  Z)   R  Q  Z
○-equiv R Q = record { left = left R Q
                     ; right = right R Q
                     ; leftright = leftright R Q
                     ; rightleft = rightleft R Q
                     }