{-# OPTIONS --safe #-}
open import Data.Nat.Base using (ℕ)
module Data.StagedList (Item : ℕ → Set) where

open import Data.Nat.Base using (ℕ; zero; suc; _≥‴_; _≤‴_; ≤‴-refl; ≤‴-step)
open import Data.Sum.Base as Sum using (_⊎_; inj₁; inj₂)
open import Function.Base using (id; _∘_)
open import Relation.Binary.PropositionalEquality.Core using (_≡_; refl; sym; trans; cong; cong₂; cong-app)
open import Relation.Nullary using (¬_)


variable
  m n m′ n′ : ℕ
  m≥n : m ≥‴ n
  m′≥n : m′ ≥‴ n

≤‴-trans : n′ ≤‴ n → n ≤‴ m → n′ ≤‴ m
≤‴-trans ≤‴-refl n≤m  = n≤m
≤‴-trans (≤‴-step sn′≤m) n≤m = ≤‴-step (≤‴-trans sn′≤m n≤m)

infixl 5 _,[_^_]
data StagedList : ℕ → Set where
  ∅ : StagedList n
  _,[_^_] : StagedList n → Item m → (m≥n : m ≥‴ n) → StagedList n

infixl 5 _++_
_++_ : StagedList n → StagedList n → StagedList n
xs ++ ∅ = xs
xs ++ (ys ,[ y ^ m≥n ]) = xs ++ ys ,[ y ^ m≥n ]

infixl 5 _↾
_↾ : StagedList n → StagedList (suc n)
∅ ↾ = ∅
xs ,[ x ^ ≤‴-refl ] ↾ = xs ↾
xs ,[ x ^ ≤‴-step m≥n ] ↾ = xs ↾ ,[ x ^ m≥n ] 

infixl 5 _↾≥_
_↾≥_ : StagedList n → m ≥‴ n → StagedList m
xs ↾≥ ≤‴-refl = xs
xs ↾≥ ≤‴-step m≥sn = xs ↾ ↾≥ m≥sn

inject₁ : StagedList (suc n) → StagedList n
inject₁ ∅ = ∅
inject₁ (xs ,[ x ^ m≥sn ]) = inject₁ xs ,[ x ^ ≤‴-step m≥sn ]

inject≥ : m ≥‴ n → StagedList m → StagedList n
inject≥ ≤‴-refl xs = xs
inject≥ (≤‴-step m≥sn) xs = inject₁ (inject≥ m≥sn xs)

private variable
  xs ys zs : StagedList n
  x x′ y z : Item n

-- Properties
↾≥-↾≥ : (m≥n : m ≥‴ n) (m′≥m : m′ ≥‴ m) → xs ↾≥ m≥n ↾≥ m′≥m ≡ xs ↾≥ (≤‴-trans m≥n m′≥m)
↾≥-↾≥ ≤‴-refl m′≥m = refl
↾≥-↾≥ (≤‴-step m≥n) m′≥m = ↾≥-↾≥ m≥n m′≥m

∅-++ : ∅ ++ xs ≡ xs
∅-++ {xs = ∅} = refl
∅-++ {xs = xs ,[ x ^ m≥n ]} = cong (_,[ x ^ m≥n ]) ∅-++

++-assoc : ∀ zs → xs ++ (ys ++ zs) ≡ (xs ++ ys) ++ zs
++-assoc ∅ = refl
++-assoc (zs ,[ z ^ m≥n ]) = cong (_,[ z ^ m≥n ]) (++-assoc zs)

++-↾ : ∀ ys → (xs ++ ys) ↾ ≡ xs ↾ ++ (ys ↾)
++-↾ ∅ = refl
++-↾ (ys ,[ y ^ ≤‴-refl ]) = ++-↾ ys
++-↾ (ys ,[ y ^ ≤‴-step m≥sn ]) = cong (_,[ y ^ m≥sn ]) (++-↾ ys)

++-↾≥ : ∀ ys (m≥n : m ≥‴ n) → (xs ++ ys) ↾≥ m≥n ≡ xs ↾≥ m≥n ++ (ys ↾≥ m≥n)
++-↾≥ ys ≤‴-refl = refl
++-↾≥ ys (≤‴-step m≥sn) = trans (cong (_↾≥ m≥sn) (++-↾ ys)) (++-↾≥ (ys ↾) m≥sn)

,-↾≥ : ∀ (m≥n : m ≥‴ n) → xs ,[ x ^ m≥n ] ↾≥ m≥n ≡ xs ↾≥ m≥n ,[ x ^ ≤‴-refl ]
,-↾≥ ≤‴-refl = refl
,-↾≥ (≤‴-step m≥n) = ,-↾≥ m≥n

inject₁-↾ : ∀ (xs : StagedList (suc n)) → (inject₁ xs) ↾ ≡ xs
inject₁-↾ ∅ = refl
inject₁-↾ (xs ,[ x ^ m≥n ]) = cong (_,[ x ^ m≥n ]) (inject₁-↾ xs)

inject₁-++ : ∀ (ys : StagedList (suc n)) → inject₁ (xs ++ ys) ≡ inject₁ xs ++ inject₁ ys
inject₁-++ ∅ = refl
inject₁-++ (ys ,[ y ^ m≥n ]) = cong (_,[ y ^ ≤‴-step m≥n ]) (inject₁-++ ys)

++-inject₁-↾ : ∀ ys → (xs ++ inject₁ ys) ↾ ≡ xs ↾ ++ ys
++-inject₁-↾ ∅ = refl
++-inject₁-↾ (ys ,[ y ^ m≥n ]) = cong (_,[ y ^ m≥n ]) (++-inject₁-↾ ys)

inject≥-∅ : (m≥n : m ≥‴ n) → inject≥ m≥n ∅ ≡ ∅
inject≥-∅ ≤‴-refl = refl
inject≥-∅ (≤‴-step m≥n) = cong inject₁ (inject≥-∅ m≥n) 

-- Membership
infix 4 _∋[_^_]
data _∋[_^_] : StagedList n → Item m → m ≥‴ n → Set where
  Z : xs ,[ x ^ m≥n ] ∋[ x ^ m≥n ]
  S : xs ∋[ x ^ m≥n ] → xs ,[ y ^ m′≥n ] ∋[ x ^ m≥n ]

-- Membership properties
∋-++⁻ : ∀ ys → xs ++ ys ∋[ x ^ m≥n ] → xs ∋[ x ^ m≥n ] ⊎ ys ∋[ x ^ m≥n ]
∋-++⁻ ∅ xs∋x = inj₁ xs∋x
∋-++⁻ (ys ,[ x ^ m≥n ]) Z = inj₂ Z
∋-++⁻ (ys ,[ x ^ m≥n ]) (S xs++ys∋x) = Sum.map id S (∋-++⁻ ys xs++ys∋x)

∋-++⁺ˡ : ∀ ys → xs ∋[ x ^ m≥n ] → xs ++ ys ∋[ x ^ m≥n ]
∋-++⁺ˡ ∅ = id
∋-++⁺ˡ (ys ,[ y ^ m≥n ]) = S ∘ ∋-++⁺ˡ ys

∋-++⁺ʳ : ∀ xs → ys ∋[ y ^ m≥n ] → xs ++ ys ∋[ y ^ m≥n ]
∋-++⁺ʳ xs Z = Z
∋-++⁺ʳ xs (S ys∋y) = S (∋-++⁺ʳ xs ys∋y)

∋-↾⁻ : ∀ {m≥sn} → xs ↾ ∋[ x ^ m≥sn ] → xs ∋[ x ^ ≤‴-step m≥sn ]
∋-↾⁻ {xs = xs ,[ x ^ ≤‴-refl ]} xs↾∋x = S (∋-↾⁻ xs↾∋x)
∋-↾⁻ {xs = xs ,[ x ^ ≤‴-step m≥n ]} Z = Z
∋-↾⁻ {xs = xs ,[ x ^ ≤‴-step m≥n ]} (S xs↾∋x) = S (∋-↾⁻ xs↾∋x)

∋-↾⁺ : ∀ {m≥sn} → xs ∋[ x ^ ≤‴-step m≥sn ] → xs ↾ ∋[ x ^ m≥sn ]
∋-↾⁺ Z = Z
∋-↾⁺ (S {m′≥n = ≤‴-refl} xs∋x) = ∋-↾⁺ xs∋x
∋-↾⁺ (S {m′≥n = ≤‴-step m′≥n} xs∋x) = S (∋-↾⁺ xs∋x)

∋-↾≥⁻ : ∀ n≥n′ → xs ↾≥ n≥n′ ∋[ x ^ m≥n ] → xs ∋[ x ^ ≤‴-trans n≥n′ m≥n ] 
∋-↾≥⁻ ≤‴-refl = id
∋-↾≥⁻ (≤‴-step n≥n′) = ∋-↾⁻ ∘ ∋-↾≥⁻ n≥n′

∋-↾≥⁺ : ∀ m≥n → xs ∋[ x ^ m≥n ] → xs ↾≥ m≥n ∋[ x ^ ≤‴-refl ]
∋-↾≥⁺ ≤‴-refl = id
∋-↾≥⁺ (≤‴-step m≥n) = ∋-↾≥⁺ m≥n ∘ ∋-↾⁺

∋-inject₁⁺ : xs ∋[ x ^ m≥n ] → inject₁ xs ∋[ x ^ ≤‴-step m≥n ]
∋-inject₁⁺ Z = Z
∋-inject₁⁺ (S xs∋x) = S (∋-inject₁⁺ xs∋x)

∋-inject≥⁺ : (n≥n′ : n ≥‴ n′) → xs ∋[ x ^ m≥n ] → inject≥ n≥n′ xs ∋[ x ^ ≤‴-trans n≥n′ m≥n ]
∋-inject≥⁺ ≤‴-refl = id
∋-inject≥⁺ (≤‴-step n≥sn′) = ∋-inject₁⁺ ∘ ∋-inject≥⁺ n≥sn′

inject₁-∌-refl : ¬ (inject₁ xs ∋[ x ^ ≤‴-refl ])
inject₁-∌-refl {xs = xs ,[ _ ^ _ ]} (S xs∋x) = inject₁-∌-refl xs∋x

∋-castˡ : xs ≡ ys → xs ∋[ x ^ m≥n ] → ys ∋[ x ^ m≥n ]
∋-castˡ refl x = x

Rename : StagedList n → StagedList n → Set
Rename xs ys = ∀ {m} {x} {m≥n : m ≥‴ _} → xs ∋[ x ^ m≥n ] → ys ∋[ x ^ m≥n ]

ext : Rename xs ys → Rename (xs ,[ z ^ m≥n ]) (ys ,[ z ^ m≥n ])
ext ρ Z = Z
ext ρ (S xs∋x) = S (ρ xs∋x)

ext++ : ∀ zs → Rename xs ys → Rename (xs ++ zs) (ys ++ zs)
ext++ ∅ ρ = ρ
ext++ (zs ,[ z ^ m≥n ]) ρ = ext (ext++ zs ρ)

ext↾ : Rename xs ys → Rename (xs ↾) (ys ↾)
ext↾ ρ = ∋-↾⁺ ∘ ρ ∘ ∋-↾⁻ 

ext↾≥ : ∀ (m≥n : m ≥‴ n) → Rename xs ys → Rename (xs ↾≥ m≥n) (ys ↾≥ m≥n)
ext↾≥ ≤‴-refl ρ = ρ
ext↾≥ (≤‴-step m≥sn) ρ = ext↾≥ m≥sn (ext↾ ρ)

ext↾≥++ : ∀ (m≥n : m ≥‴ n) zs → Rename xs ys → Rename (xs ↾≥ m≥n ++ zs) (ys ↾≥ m≥n ++ zs)
ext↾≥++ m≥n zs = ext++ zs ∘ ext↾≥ m≥n

↾-inject₁ : (xs : StagedList n) → Rename (inject₁ (xs ↾)) xs
↾-inject₁ (xs ,[ x ^ ≤‴-refl ]) = S ∘ ↾-inject₁ xs
↾-inject₁ (xs ,[ x ^ ≤‴-step m≥n ]) Z = Z
↾-inject₁ (xs ,[ x ^ ≤‴-step m≥n ]) (S xs∋x) = S (↾-inject₁ xs xs∋x)

_≗_ : Rename xs ys → Rename xs ys → Set
_≗_ {xs = xs} ρ ρ′ = ∀ {m} {x} {m≥n : m ≥‴ _} (xs∋x : xs ∋[ x ^ m≥n ]) → ρ xs∋x ≡ ρ′ xs∋x

-- Renaming properties
≗-refl : {ρ : Rename xs ys} → ρ ≗ ρ
≗-refl _ = refl

≗-trans : {ρ ρ′ ρ″ : Rename xs ys} → ρ ≗ ρ′ → ρ′ ≗ ρ″ → ρ ≗ ρ″
≗-trans ρ≗ρ′ ρ′≗ρ″ xs∋x = trans (ρ≗ρ′ xs∋x) (ρ′≗ρ″ xs∋x)

ext-resp-≗ : {ρ ρ′ : Rename xs ys} → ρ ≗ ρ′ → ext {z = z} {m≥n = m≥n} ρ ≗ ext ρ′
ext-resp-≗ ρ≗ρ′ Z = refl
ext-resp-≗ ρ≗ρ′ (S _) = cong S (ρ≗ρ′ _)

ext-id : ext {xs = xs} {z = z} {m≥n = m≥n} id ≗ id
ext-id Z = refl
ext-id (S _) = refl

ext++-id : ∀ (zs : StagedList n) → ext++ {xs = xs} zs id ≗ id
ext++-id ∅ = ≗-refl
ext++-id (zs ,[ z ^ m≥n ]) = ≗-trans (ext-resp-≗ (ext++-id zs)) ext-id

{-
infix 5 _,_
data All {n} (P : ∀ {m} → Item m → m ≥‴ n → Set) : StagedList n → Set where
  ∅ : All P ∅
  _,_ : All P xs → P x m≥n → All P (xs ,[ x ^ m≥n ])

private module _ {P : ∀ {m} → Item m → m ≥‴ n → Set} where
  lookup : All P xs → xs ∋[ x ^ m≥n ] → P x m≥n
  lookup (σ , e) Z = e         
  lookup (σ , _) (S x) = lookup σ x
-}