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