module Ornament.ParallelComposition.PredicateSwap where

open import Prelude
open import Refinement
open import Description
open import Ornament
open import Ornament.ParallelComposition
open import Ornament.RefinementSemantics

open import Function using (id; _∘_)
open import Data.Unit using (; tt)
open import Data.Product using (Σ; _,_; proj₁; proj₂; _×_; <_,_>; uncurry) renaming (map to _**_)
open import Relation.Binary.PropositionalEquality using (_≡_; refl; subst; cong; cong₂)


--------
-- projection

project-Ind :  {I J K} {e : J  I} {f : K  I} {D E F} (O : Orn e D E) (P : Orn f D F)   {i}  μ D i  Set
project-Ind {e = e} {f} O P {i} x = {j : e ⁻¹ i} {k : f ⁻¹ i}  [ ok (j , k) ] x   O  P   [ j ] x  O × [ k ] x  P

project-aux :
   {I J K} {e : J  I} {f : K  I} {D E F} {O : Orn e D E} {P : Orn f D F} 
   {D' E' F'} (O' : ROrn e D' E') (P' : ROrn f D' F') (xs :  D'  (μ D)) (ih : All D' (project-Ind O P) xs) 
   toRDesc (pcROrn (toROrn (pcROrn O' P')) (toROrn (erode D' xs)))  (μ (cpD  O  P )) 
   toRDesc (pcROrn O' (toROrn (erode D' xs)))  (μ (cpD O)) ×  toRDesc (pcROrn P' (toROrn (erode D' xs)))  (μ (cpD P))
project-aux                     xs         ih         rs                 = tt , tt
project-aux           (Δ T P')   xs         ih         (t , rs)           = (id ** (_,_ t)) (project-aux  (P' t) xs ih rs)
project-aux (ṿ idx)    (ṿ idx')   xs         ih         r                  = ih r
project-aux (ṿ idx)    (Δ T P')   xs         ih         (t , rs)           = (id ** (_,_ t)) (project-aux (ṿ idx) (P' t) xs ih rs)
project-aux (σ S O')   (σ .S P')  (s , xs)   ih         rs                 = project-aux (O' s) (P' s) xs ih rs
project-aux (σ S O')   (Δ T P')   (s , xs)   ih         (t , rs)           = (id ** (_,_ t)) (project-aux (σ S O') (P' t) (s , xs) ih rs)
project-aux (σ S O')   ( s P')   (.s , xs)  ih         (refl , rs)        = (id ** (_,_ refl)) (project-aux (O' s) P' xs ih rs)
project-aux (Δ T O')   P'         xs         ih         (t , rs)           = ((_,_ t) ** id) (project-aux (O' t) P' xs ih rs)
project-aux ( s O')   (σ S P')   (.s , xs)  ih         (refl , rs)        = ((_,_ refl) ** id) (project-aux O' (P' s) xs ih rs)
project-aux ( s O')   (Δ T P')   (s' , xs)  ih         (t , rs)           = (id ** (_,_ t)) (project-aux ( s O') (P' t) (s' , xs) ih rs)
project-aux ( s O')   ( .s P')  (.s , xs)  ih         (refl , refl , rs) = ((_,_ refl) ** (_,_ refl)) (project-aux O' P' xs ih rs)
project-aux (O' * O'') (Δ T P')   xs         ih         (t , rs)           = (id ** (_,_ t)) (project-aux (O' * O'') (P' t) xs ih rs)
project-aux (O' * O'') (P' * P'') (xs , xs') (ih , ih') (rs , rs')         = let x = project-aux O' P' xs ih rs
                                                                                 y = project-aux O'' P'' xs' ih' rs'
                                                                             in (proj₁ x , proj₁ y) , (proj₂ x , proj₂ y)

project-ind :
   {I J K} {e : J  I} {f : K  I} {D E F} (O : Orn e D E) (P : Orn f D F) 
  {i : I} (xs :  D (μ D) i)  All (D at i) (project-Ind O P) xs  project-Ind O P (con xs)
project-ind {e = e} {f} {D} {E} {F} (wrap O) (wrap P) xs ih {j} {k} (con rs) =
  (con ** con) (project-aux (subst  i  ROrn e (D at i) (E at und j)) (to≡ j) (O (und j)))
                            (subst  i  ROrn f (D at i) (F at und k)) (to≡ k) (P (und k))) xs ih rs)

project :  {I J K} {e : J  I} {f : K  I} {D E F} (O : Orn e D E) (P : Orn f D F) 
             {i} (x : μ D i) {j : e ⁻¹ i} {k : f ⁻¹ i}  [ ok (j , k) ] x   O  P   [ j ] x  O × [ k ] x  P
project {e = e} {f} {D} {E} {F} O P x r = induction D (project-Ind O P) (project-ind O P) x r

project-injective :
   {I J K} {e : J  I} {f : K  I} {D E F} (O : Orn e D E) (P : Orn f D F) 
   {i} (x : μ D i) {j : e ⁻¹ i} {k : f ⁻¹ i} (r r' : [ ok (j , k) ] x   O  P ) 
  project O P x r  project O P x r'  r  r'
project-injective {e = e} {f} {D} {E} {F} (wrap O!) (wrap P!) x r r' eq =
  induction D Ind  { xs ih {j} {k} (con rs) (con rs') eq 
                       cong con (aux (subst  i  ROrn e (D at i) (E at und j)) (to≡ j) (O! (und j)))
                                             (subst  i  ROrn f (D at i) (F at und k)) (to≡ k) (P! (und k)))
                                             xs ih rs rs' (cong (decon ** decon) eq))}) x r r' eq
  where
    O : Orn e D E
    O = wrap {E = E} O!
    P : Orn f D F
    P = wrap {E = F} P!
    Ind :  {i}  μ D i  Set
    Ind {i} x = {j : e ⁻¹ i} {k : f ⁻¹ i} (r r' : [ ok (j , k) ] x   O  P )  project O P x r  project O P x r'  r  r'
    aux :  {D' E' F'} (O' : ROrn e D' E') (P' : ROrn f D' F') (xs :  D'  (μ D)) (ih : All D' Ind xs) 
          (rs rs' :  toRDesc (pcROrn (toROrn (pcROrn O' P')) (toROrn (erode D' xs)))  (μ (cpD  O  P ))) 
          project-aux O' P' xs (everywhereInduction D D' (project-Ind O P) (project-ind O P) xs) rs
             project-aux O' P' xs (everywhereInduction D D' (project-Ind O P) (project-ind O P) xs) rs'  rs  rs'
    aux   xs ih rs rs' eq = refl
    aux  (Δ T P') xs ih (t , rs) (t' , rs') eq =
      subst  t'  (rs' :  toRDesc (pcROrn (toROrn (pcROrn  (P' t'))) )  (μ (cpD  O  P )))
                    (eq : let f = λ t rs  project-aux  (P' t) xs (everywhereInduction D  (project-Ind O P) (project-ind O P) xs) rs
                          in (proj₁ (f t rs) , t , proj₂ (f t rs))  (proj₁ (f t' rs') , t' , proj₂ (f t' rs')))  (t , rs)  (t' , rs'))
            (cong proj₁ (cong proj₂ eq))
             rs' eq  cong (_,_ t) (aux  (P' t) xs ih rs rs' (cong₂ _,_ (cong proj₁ eq) (cong-proj₂ (cong proj₂ eq))))) rs' eq
    aux (ṿ idx) (ṿ idx') xs ih rs rs' eq = ih rs rs' eq
    aux {ṿ i} (ṿ idx) (Δ T P') xs ih (t , rs) (t' , rs') eq =
      subst  t'  (rs' :  toRDesc (pcROrn (toROrn (pcROrn (ṿ idx) (P' t'))) (ṿ refl))  (μ (cpD  O  P )))
                    (eq : let f = λ t rs  project-aux (ṿ idx) (P' t) xs (everywhereInduction D (ṿ i) (project-Ind O P) (project-ind O P) xs) rs
                          in (proj₁ (f t rs) , t , proj₂ (f t rs))  (proj₁ (f t' rs') , t' , proj₂ (f t' rs')))  (t , rs)  (t' , rs'))
            (cong proj₁ (cong proj₂ eq))
             rs' eq  cong (_,_ t) (aux (ṿ idx) (P' t) xs ih rs rs' (cong₂ _,_ (cong proj₁ eq) (cong-proj₂ (cong proj₂ eq))))) rs' eq
    aux (σ S O') (σ .S P') (s , xs) ih rs rs' eq = aux (O' s) (P' s) xs ih rs rs' eq
    aux {σ .S D'} (σ S O') (Δ T P') (s , xs) ih (t , rs) (t' , rs') eq =
      subst  t'  (rs' :  toRDesc (pcROrn (toROrn (pcROrn (σ S O') (P' t'))) ( s (toROrn (erode (D' s) xs))))  (μ (cpD  O  P )))
                    (eq : let f = λ t rs  project-aux (σ S O') (P' t) (s , xs)
                                                       (everywhereInduction D (σ S D') (project-Ind O P) (project-ind O P) (s , xs)) rs
                          in (proj₁ (f t rs) , t , proj₂ (f t rs))  (proj₁ (f t' rs') , t' , proj₂ (f t' rs')))  (t , rs)  (t' , rs'))
            (cong proj₁ (cong proj₂ eq))
             rs' eq  cong (_,_ t) (aux (σ S O') (P' t) (s , xs) ih rs rs' (cong₂ _,_ (cong proj₁ eq) (cong-proj₂ (cong proj₂ eq))))) rs' eq
    aux (σ S O') ( s P') (.s , xs) ih (refl , rs) (refl , rs') eq =
      cong (_,_ refl) (aux (O' s) P' xs ih rs rs' (cong₂ _,_ (cong proj₁ eq) (cong-proj₂ (cong proj₂ eq))))
    aux {D'} (Δ T O') P' xs ih (t , rs) (t' , rs') eq =
      subst  t'  (rs' :  toRDesc (pcROrn (toROrn (pcROrn (O' t') P')) (toROrn (erode D' xs)))  (μ (cpD  O  P )))
                    (eq : let f = λ t rs  project-aux (O' t) P' xs
                                                       (everywhereInduction D D' (project-Ind O P) (project-ind O P) xs) rs
                          in ((t , proj₁ (f t rs)) , proj₂ (f t rs))  ((t' , proj₁ (f t' rs')) , proj₂ (f t' rs')))  (t , rs)  (t' , rs'))
            (cong proj₁ (cong proj₁ eq))
             rs' eq  cong (_,_ t) (aux (O' t) P' xs ih rs rs' (cong₂ _,_ (cong-proj₂ (cong proj₁ eq)) (cong proj₂ eq)))) rs' eq
    aux ( s O') (σ S P') (.s , xs) ih (refl , rs) (refl , rs') eq =
      cong (_,_ refl) (aux O' (P' s) xs ih rs rs' (cong₂ _,_ (cong-proj₂ (cong proj₁ eq)) (cong proj₂ eq)))
    aux {σ S D'} ( s O') (Δ T P') (s' , xs) ih (t , rs) (t' , rs') eq =
      subst  t'  (rs' :  toRDesc (pcROrn (toROrn (pcROrn ( s O') (P' t'))) ( s' (toROrn (erode (D' s') xs))))  (μ (cpD  O  P )))
                    (eq : let f = λ t rs  project-aux ( s O') (P' t) (s' , xs)
                                                       (everywhereInduction D (σ S D') (project-Ind O P) (project-ind O P) (s' , xs)) rs
                          in (proj₁ (f t rs) , t , proj₂ (f t rs))  (proj₁ (f t' rs') , t' , proj₂ (f t' rs')))  (t , rs)  (t' , rs'))
            (cong proj₁ (cong proj₂ eq))
             rs' eq  cong (_,_ t) (aux ( s O') (P' t) (s' , xs) ih rs rs' (cong₂ _,_ (cong proj₁ eq) (cong-proj₂ (cong proj₂ eq))))) rs' eq
    aux ( s O') ( .s P') (.s , xs) ih (refl , refl , rs) (refl , refl , rs') eq =
      cong (_,_ refl) (cong (_,_ refl) (aux O' P' xs ih rs rs' (cong₂ _,_ (cong-proj₂ (cong proj₁ eq)) (cong-proj₂ (cong proj₂ eq)))))
    aux {D' * D''} (O' * O'') (Δ T P') (xs , xs') ih (t , rs) (t' , rs') eq =
      subst  t'  (rs' :  toRDesc (pcROrn (toROrn (pcROrn (O' * O'') (P' t'))) (toROrn (erode (D' * D'') (xs , xs'))))  (μ (cpD  O  P )))
                    (eq : let f = λ t rs  project-aux (O' * O'') (P' t) (xs , xs')
                                                       (everywhereInduction D (D' * D'') (project-Ind O P) (project-ind O P) (xs , xs')) rs
                          in (proj₁ (f t rs) , t , proj₂ (f t rs))  (proj₁ (f t' rs') , t' , proj₂ (f t' rs')))  (t , rs)  (t' , rs'))
            (cong proj₁ (cong proj₂ eq))
             rs' eq  cong (_,_ t) (aux (O' * O'') (P' t) (xs , xs') ih rs rs' (cong₂ _,_ (cong proj₁ eq) (cong-proj₂ (cong proj₂ eq))))) rs' eq
    aux (O' * O'') (P' * P'') (xs , xs') (ih , ih') (rs , rs') (rs'' , rs''') eq =
      cong₂ _,_ (aux O' P' xs ih rs rs'' (cong (proj₁ ** proj₁) eq)) (aux O'' P'' xs' ih' rs' rs''' (cong (proj₂ ** proj₂) eq))


--------
-- integration

integrate-Ind :  {I J K} {e : J  I} {f : K  I} {D E F} (O : Orn e D E) (P : Orn f D F)   {i}  μ D i  Set
integrate-Ind {e = e} {f} O P {i} x =
  {j : e ⁻¹ i} {k : f ⁻¹ i}  (r : [ j ] x  O) (s : [ k ] x  P)  Σ[ t  [ ok (j , k) ] x   O  P  ] project O P x t  (r , s)

integrate-aux :
   {I J K} {e : J  I} {f : K  I} {D E F} {O : Orn e D E} {P : Orn f D F}  
   {D' E' F'} (O' : ROrn e D' E') (P' : ROrn f D' F') (xs :  D'  (μ D)) (ih : All D' (integrate-Ind O P) xs) 
  (rs :  toRDesc (pcROrn O' (toROrn (erode D' xs)))  (μ (cpD O))) (ss :  toRDesc (pcROrn P' (toROrn (erode D' xs)))  (μ (cpD P))) 
  Σ[ ts   toRDesc (pcROrn (toROrn (pcROrn O' P')) (toROrn (erode D' xs)))  (μ (cpD  O  P )) ]
    project-aux O' P' xs (everywhereInduction D D' (project-Ind O P) (project-ind O P) xs) ts  (rs , ss)
integrate-aux                     xs         ih         rs          ss          = tt , refl
integrate-aux           (Δ T O')   xs         ih         rs          (t , ss)    = ((_,_ t) ** cong (id ** (_,_ t)))
                                                                                      (integrate-aux  (O' t) xs ih rs ss)
integrate-aux (ṿ idx)    (ṿ idx')   xs         ih         r           s           = ih r s
integrate-aux (ṿ idx)    (Δ T P')   xs         ih         rs          (t , ss)    = ((_,_ t) ** cong (id ** (_,_ t)))
                                                                                      (integrate-aux (ṿ idx) (P' t) xs ih rs ss)
integrate-aux (σ S O')   (σ .S P')  (s , xs)   ih         rs          ss          = integrate-aux (O' s) (P' s) xs ih rs ss
integrate-aux (σ S O')   (Δ T P')   (s , xs)   ih         rs          (t , ss)    = ((_,_ t) ** cong (id ** (_,_ t)))
                                                                                      (integrate-aux (σ S O') (P' t) (s , xs) ih rs ss)
integrate-aux (σ S O')   ( s P')   (.s , xs)  ih         rs          (refl , ss) = ((_,_ refl) ** cong (id ** (_,_ refl)))
                                                                                      (integrate-aux (O' s) P' xs ih rs ss)
integrate-aux (Δ T O')   P'         xs         ih         (t , rs)    ss          = ((_,_ t) ** cong ((_,_ t) ** id))
                                                                                      (integrate-aux (O' t) P' xs ih rs ss)
integrate-aux ( s O')   (σ S P')   (.s , xs)  ih         (refl , rs) ss          = ((_,_ refl) ** cong ((_,_ refl) ** id))
                                                                                      (integrate-aux O' (P' s) xs ih rs ss)
integrate-aux ( s O')   (Δ T P')   (s' , xs)  ih         rs          (t , ss)    = ((_,_ t) ** cong (id ** (_,_ t)))
                                                                                      (integrate-aux ( s O') (P' t) (s' , xs) ih rs ss)
integrate-aux ( s O')   ( .s P')  (.s , xs)  ih         (refl , rs) (refl , ss) = ((λ ts  refl , refl , ts) ** cong (_,_ refl ** _,_ refl))
                                                                                      (integrate-aux O' P' xs ih rs ss)
integrate-aux (O' * O'') (Δ T P')   xs         ih         rs          (t , ss)    = (_,_ t ** cong (id ** _,_ t))
                                                                                      (integrate-aux (O' * O'') (P' t) xs ih rs ss)
integrate-aux (O' * O'') (P' * P'') (xs , xs') (ih , ih') (rs , rs')  (ss , ss')  =
  let a = integrate-aux O' P' xs ih rs ss
      b = integrate-aux O'' P'' xs' ih' rs' ss'
  in (proj₁ a , proj₁ b) , cong₂ _,_ (cong₂ _,_ (cong proj₁ (proj₂ a)) (cong proj₁ (proj₂ b)))
                                    (cong₂ _,_ (cong proj₂ (proj₂ a)) (cong proj₂ (proj₂ b)))

integrate :  {I J K} {e : J  I} {f : K  I} {D E F} (O : Orn e D E) (P : Orn f D F) 
             {i} (x : μ D i) {j : e ⁻¹ i} {k : f ⁻¹ i} 
            (r : [ j ] x  O) (s : [ k ] x  P)  Σ[ t  [ ok (j , k) ] x   O  P  ] project O P x t  (r , s)
integrate {e = e} {f} {D} {E} {F} (wrap O) (wrap P) x r s =
  induction D (integrate-Ind (wrap O) (wrap P))
     { xs ih {j} {k} (con rs) (con ss) 
         (con ** cong (con ** con)) (integrate-aux (subst  i  ROrn e (D at i) (E at und j)) (to≡ j) (O (und j)))
                                                   (subst  i  ROrn f (D at i) (F at und k)) (to≡ k) (P (und k))) xs ih rs ss) }) x r s


--------
-- predicate swap

Swap-⊗ :  {I J K} {e : J  I} {f : K  I} {D E F} (O : Orn e D E) (P : Orn f D F) 
         Swap (RSem O)  Swap (RSem P)  Swap (RSem  O  P )
Swap-⊗ O P s t =
  record { Q = λ { {._} (ok (j , k)) x  Swap.Q s j x × Swap.Q t k x }
         ; s = λ { {._} (ok (j , k)) x  transIso (toIso (project O P x) (project-injective O P x) (uncurry (integrate O P x)))
                                                  (prodIso (Swap.s s j x) (Swap.s t k x)) } }