module Ornament.RefinementSemantics where

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

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


--------
-- canonical predicate

cpD : ∀ {I J} {e : J → I} {D E} (O : Orn e D E) → Desc (e ⋈ proj₁)
cpD {D = D} O = ⌊ O ⊗ ⌈ singOrn D ⌉ ⌋

cpOrn : ∀ {I J} {e : J → I} {D E} (O : Orn e D E) → Orn π₁ E (cpD O)
cpOrn {D = D} O = diffOrn-l O ⌈ singOrn D ⌉

[_]_⊩_ : ∀ {I J} {e : J → I} {D E} → ∀ {i} → e ⁻¹ i → (x : μ D i) → Orn e D E → Set
[ j ] x ⊩ O = μ (cpD O) (j , (ok (_ , x)))

remember-aux :
  ∀ {I J} {e : J → I} {D E} {O : Orn e D E} →
  ∀ {D' E'} (O' : ROrn e D' E') (ys : ⟦ E' ⟧ (μ E)) (ih : All E' (λ {j} y → [ ok j ] forget O y ⊩ O) ys) →
  ⟦ toRDesc (pcROrn O' (toROrn (erode D' (erase O' (mapFold E E' {μ D ∘ e} (λ {i} y → con (erase (unwrap O i) y)) ys))))) ⟧ (μ (cpD O))
remember-aux              ∎          ys         ih         = tt
remember-aux {O = wrap _} (ṿ refl)   ys         ih         = ih
remember-aux              (σ S O')   (s , ys)   ih         = remember-aux (O' s) ys ih
remember-aux              (Δ T O')   (t , ys)   ih         = t , remember-aux (O' t) ys ih
remember-aux              (∇ s O')   ys         ih         = refl , remember-aux O' ys ih
remember-aux              (O' * O'') (ys , ys') (ih , ih') = remember-aux O' ys ih , remember-aux O'' ys' ih'

remember : ∀ {I J} {e : J → I} {D E} (O : Orn e D E) → ∀ {j} (y : μ E j) → [ ok j ] forget O y ⊩ O
remember {e = e} {D} {E} (wrap O) =
  induction E (λ {j} y → [ ok j ] forget (wrap O) y ⊩ wrap O) (λ { {j} ys ih → con (remember-aux (O j) ys ih)})

remember-forget-inverse :
  ∀ {I J} {e : J → I} {D E} (O : Orn e D E) → ∀ {j} {r : Σ[ x ∶ μ D (e j) ] [ ok j ] x ⊩ O} →
  <_,_> {C = λ x → [ ok j ] x ⊩ O} (forget O) (remember O) (forget (cpOrn O) (proj₂ r)) ≡ r
remember-forget-inverse {e = e} {D} {E} (wrap O!) = induction D Ind aux _ refl _
  where
    O : Orn e D E
    O = wrap {E = E} O!
    Ind' : ∀ {j i} → e j ≡ i → μ D i → Set
    Ind' {j} refl x = (r : [ ok j ] x ⊩ O) → <_,_> {C = λ x → [ ok j ] x ⊩ O} (forget O) (remember O) (forget (cpOrn O) r) ≡ (x , r)
    Ind : ∀ {i} → μ D i → Set
    Ind {i} x = ∀ {j} → (eq : e j ≡ i) → Ind' eq x
    aux' : ∀ {D' E'} (O' : ROrn e D' E') (xs : ⟦ D' ⟧ (μ D)) (ih : All D' Ind xs) (rs : ⟦ toRDesc (pcROrn O' (toROrn (erode D' xs))) ⟧ (μ (cpD O))) → _≡_ {A = Σ[ xs ∶ ⟦ D' ⟧ (μ D) ] ⟦ toRDesc (pcROrn O' (toROrn (erode D' xs))) ⟧ (μ (cpD O))} (erase O' (mapFold E E' (ornAlg O) (erase (diffROrn-l O' (toROrn (erode D' xs))) (mapFold (cpD O) (toRDesc (pcROrn O' (toROrn (erode D' xs)))) {μ E ∘ π₁} (λ {i} → ornAlg (cpOrn O) {i}) rs))) , remember-aux O' (erase (diffROrn-l O' (toROrn (erode D' xs))) (mapFold (cpD O) (toRDesc (pcROrn O' (toROrn (erode D' xs)))) (λ {i} → ornAlg (cpOrn O) {i}) rs)) (everywhereInduction E E' (λ {j} x → [ ok j ] forget O x ⊩ O) (λ {i} xs ih → con (remember-aux (O! i) xs ih)) (erase (diffROrn-l O' (toROrn (erode D' xs))) (mapFold (cpD O) (toRDesc (pcROrn O' (toROrn (erode D' xs)))) (λ {i} → ornAlg (cpOrn O) {i}) rs)))) (xs , rs)
    aux' ∎          xs         ih         rs          = refl
    aux' (ṿ refl)   xs         ih         r           = ih refl r
    aux' (σ S O')   (s , xs)   ih         rs          = cong ((_,_ s) ** id) (aux' (O' s) xs ih rs)
    aux' (Δ T O')   xs         ih         (t , rs)    = cong (id ** (_,_ t)) (aux' (O' t) xs ih rs)
    aux' (∇ s O')   (.s , xs)  ih         (refl , rs) = cong ((_,_ s) ** (_,_ refl)) (aux' O' xs ih rs)
    aux' (O' * O'') (xs , xs') (ih , ih') (rs , rs')  =
      cong₂ (λ xsrs xs'rs' → (proj₁ xsrs , proj₁ xs'rs') , (proj₂ xsrs , proj₂ xs'rs')) (aux' O' xs ih rs) (aux' O'' xs' ih' rs')
    aux : ∀ {i} (xs : ⟦ D at i ⟧ (μ D)) (ih : All (D at i) Ind xs) → ∀ {j} → (eq : e j ≡ i) → Ind' eq (con xs)
    aux xs ih {j} refl = λ { (con rs) → cong (con ** con) (aux' (O! j) xs ih rs) }

forget-remember-inverse : ∀ {I J} {e : J → I} {D E} (O : Orn e D E) → ∀ {j} {x : μ E j} → forget (cpOrn O) (remember O x) ≡ x
forget-remember-inverse {e = e} {D} {E} (wrap O!) =
  induction E (λ x → forget (cpOrn O) (remember O x) ≡ x) (λ {j} ys ih → cong con (aux (O! j) ys ih)) _
  where
    O : Orn e D E
    O = wrap {E = E} O!
    aux : ∀ {D' E'} (O' : ROrn e D' E') (ys : ⟦ E' ⟧ (μ E)) (ih : All E' (λ x → forget (cpOrn O) (remember O x) ≡ x) ys) → erase (diffROrn-l O' (toROrn (erode D' (erase O' (mapFold E E' (ornAlg O) ys))))) (mapFold (cpD O) (toRDesc (pcROrn O' (toROrn (erode D' (erase O' (mapFold E E' (ornAlg O) ys)))))) {μ E ∘ π₁} (λ {i} → ornAlg (cpOrn O) {i}) (remember-aux O' ys (everywhereInduction E E' (λ {j} x → [ ok j ] forget O x ⊩ O) (λ {i} xs ih → con (remember-aux (O! i) xs ih)) ys))) ≡ ys
    aux ∎          ys         ih         = refl
    aux (ṿ refl)   ys         ih         = ih
    aux (σ S O')   (s , ys)   ih         = cong (_,_ s) (aux (O' s) ys ih)
    aux (Δ T O')   (t , ys)   ih         = cong (_,_ t) (aux (O' t) ys ih)
    aux (∇ s O')   ys         ih         = aux O' ys ih
    aux (O' * O'') (ys , ys') (ih , ih') = cong₂ _,_ (aux O' ys ih) (aux O'' ys' ih')


--------
-- refinement semantics

RSem : ∀ {I J} {e : J → I} {D E} → Orn e D E → Refinement (μ D) (μ E)
RSem {e = e} O =
  record { e = e
         ; P = λ j d → [ j ] d ⊩ O
         ; ℜ = λ { {._} (ok j) →
                 record { to   = < forget O , remember O >
                        ; from = forget (cpOrn O) ∘ proj₂
                        ; to-from-inverse = remember-forget-inverse O
                        ; from-to-inverse = forget-remember-inverse O } } }