open import Level
open import Function

open import Data.Unit hiding (_≟_)
open import Data.Product

open import Relation.Binary.PropositionalEquality

open import Chapter2.Logic

open import Chapter5.IDesc
open import Chapter5.IDesc.Fixpoint
open import Chapter5.IDesc.InitialAlgebra

open import Chapter8.Ornament



module Chapter8.Reornament 
         {ℓ : Level}
         {K I : Set ℓ}
         {D : func ℓ I I}
         {u : K → I}
         (o : orn D u u) 
       where

{-

open Ornament.OrnamentalAlgebra D f O

-- * 

-}

forgetIdx : Σ K (μ D ∘ u) → K
forgetIdx = proj₁ 

Extension : ∀{D}{X : I → Set ℓ} → Orn u D → ⟦ D ⟧ X → Set ℓ
Extension `1 (lift tt) = Lift ⊤
Extension (O₁ `× O₂) (xs₁ , xs₂) = Extension O₁ xs₁ × Extension O₂ xs₂
Extension (`σ T⁺) (k , xs) = Extension (T⁺ k) xs
Extension (`Σ T⁺) (s , xs) = Extension (T⁺ s) xs
Extension (`Π T⁺) f = (s : _) → Extension (T⁺ s) (f s)
Extension (`var (inv k)) xs = Lift ⊤
Extension (insert S D⁺) xs = Σ[ s ∈ S ] Extension (D⁺ s) xs
Extension {X = X} (deleteΣ {T = T} s o) (s' , xs) = Σ[ q ∈ s' ≡ s ] Extension o (subst (λ s → ⟦ T s ⟧ X) q xs) 
Extension {X = X} (deleteσ {T = T} k o) (k' , xs) = Σ[ q ∈ k' ≡ k ] Extension o (subst (λ k → ⟦ T k ⟧ X) q xs)


Structure : ∀{D'}(O : Orn u D') →
            (xs' : ⟦ D' ⟧ (μ D)) → Extension O xs' → Orn forgetIdx ⟦ O ⟧Orn
Structure `1 (lift tt) (lift tt) = `1
Structure (O₁ `× O₂) (xs₁ , xs₂) (e₁ , e₂) = Structure O₁ xs₁ e₁ `× Structure O₂ xs₂ e₂
Structure (`σ T⁺) (k , xs) e = deleteσ k (Structure (T⁺ k) xs e)
Structure (`Σ T⁺) (s , xs) e = deleteΣ s (Structure (T⁺ s) xs e)
Structure (`Π T⁺) f e = `Π (λ s → Structure (T⁺ s) (f s) (e s))
Structure (`var (inv k)) xs (lift tt) = `var (inv (k , xs))
Structure (insert S D⁺) xs (s , e) = deleteΣ s (Structure (D⁺ s) xs e)
Structure (deleteΣ s O) (.s , xs) (refl , e) = Structure O xs e
Structure (deleteσ k O) (.k , xs) (refl , e) = Structure O xs e

-- Or, merging the two:
Reorn : ∀{D'} → (O : Orn u D') → ⟦ D' ⟧ (μ D) → Orn forgetIdx ⟦ O ⟧Orn
Reorn `1 (lift tt) = `1
Reorn (O₁ `× O₂) (xs₁ , xs₂) = Reorn O₁ xs₁ `× Reorn O₂ xs₂
Reorn (`σ T⁺) (k , xs) = deleteσ k (Reorn (T⁺ k) xs)
Reorn (`Σ T⁺) (s , xs) = deleteΣ s (Reorn (T⁺ s) xs)
Reorn (`Π T⁺) f = `Π λ s → Reorn (T⁺ s) (f s)
Reorn (`var (inv k)) xs = `var (inv (k , xs))
Reorn (insert S D⁺) xs = `Σ λ s → Reorn (D⁺ s) xs
Reorn (deleteΣ {T = T} s o) (s' , xs) = insert (s' ≡ s) λ q → Reorn o (subst (λ s → ⟦ T s ⟧ (μ D)) q xs) 
Reorn (deleteσ {T = T} k o) (k' , xs) = insert (Lift (k' ≡ k)) λ { (lift q) → Reorn o (subst (λ k → ⟦ T k ⟧ (μ D)) q xs) }


reorn : orn ⟦ o ⟧orn forgetIdx forgetIdx
reorn = orn.mk (λ { (k , ⟨ xs ⟩) →
               let O = orn.out o k in
               insert (Extension O xs) (λ e → 
               Structure O xs e) })

reornD : func ℓ (Σ K (μ D ∘ u)) (Σ K (μ D ∘ u))
reornD = ⟦ reorn ⟧orn