{-# OPTIONS --without-K #-}

module tower where

open import Data.Nat
open import Data.Maybe
open import Data.Sum
open import Data.Product

open import mctt
open import extract

-- whether a context is all kinds, or else find the first type.

data AllKinds : Cxt → Set where
  AKε : AllKinds ε
  AK,,k : {cxt : Cxt} → AllKinds cxt → (k : Kind cxt) → AllKinds (cxt ,, k)

data FirstTy : Cxt → Set where
  FT,t₁ : {cxt : Cxt} → AllKinds cxt → (t : Ty cxt) → FirstTy (cxt , t)
  FT,t₊ : {cxt : Cxt} → FirstTy cxt → (t : Ty cxt) → FirstTy (cxt , t)
  FT,,k : {cxt : Cxt} → FirstTy cxt → (k : Kind cxt) → FirstTy (cxt ,, k)

AK⊎FT : (cxt : Cxt) → AllKinds cxt ⊎ FirstTy cxt
AK⊎FT ε = inj₁ AKε
AK⊎FT (cxt , t) = helper (AK⊎FT cxt) where
  helper : AllKinds cxt ⊎ FirstTy cxt → AllKinds (cxt , t) ⊎ FirstTy (cxt , t)
  helper (inj₁ ak) = inj₂ (FT,t₁ ak t)
  helper (inj₂ ft) = inj₂ (FT,t₊ ft t)
AK⊎FT (cxt ,, k) =  helper (AK⊎FT cxt) where
  helper : AllKinds cxt ⊎ FirstTy cxt → AllKinds (cxt ,, k) ⊎ FirstTy (cxt ,, k)
  helper (inj₁ ak) = inj₁ (AK,,k ak k)
  helper (inj₂ ft) = inj₂ (FT,,k ft k)

----------------------------------------
-- Extending by a single type variable
----------------------------------------

-- Insert a new kind variable X:U at the beginning, and a new type
-- variable x:X just before the first existing type variable (or at
-- the end, if there are none).  Furthermore, all kind variables
-- occurring in between are made indexed over X.
tcxt : Cxt → Cxt

-- Helper function for tcxt.
tcxt? : (cxt : Cxt) → AllKinds cxt ⊎ FirstTy cxt → Cxt

-- tcxt is easy to define in terms of its helper function.
tcxt cxt = tcxt? cxt (AK⊎FT cxt)

-- Given a context consisting of all kinds, do as tcxt but omit x:X at
-- the end.
tcxt-ak : (cxt : Cxt) → AllKinds cxt → Cxt

-- The type of the variable x:X to be added to tcxt-ak to produce tcxt.
tnewty-ak : (cxt : Cxt) → (ak : AllKinds cxt) → Ty (tcxt-ak cxt ak)

-- Given a context *not* consisting of all kinds, do as tcxt.
tcxt-ft : (cxt : Cxt) → FirstTy cxt → Cxt

-- Now tcxt? is easy to define in terms of its helper functions.
tcxt? cxt (inj₁ ak) =  (tcxt-ak cxt ak , tnewty-ak cxt ak)
tcxt? cxt (inj₂ ft) =  (tcxt-ft cxt ft)

-- Given a kind in context, find the same kind in the extended context.
tkind : {cxt : Cxt} → (k : Kind cxt) → Kind (tcxt cxt)

-- Helper function for tkind
tkind? : {cxt : Cxt} → (akft : AllKinds cxt ⊎ FirstTy cxt) → (k : Kind cxt) → Kind (tcxt? cxt akft)

-- tkind is easy to define in terms of its helper function.
tkind {cxt} k = tkind? (AK⊎FT cxt) k

-- Same for a type in context
tty : {cxt : Cxt} → (t : Ty cxt) → Ty (tcxt cxt)
tty? : {cxt : Cxt} → (akft : AllKinds cxt ⊎ FirstTy cxt) → (t : Ty cxt) → Ty (tcxt? cxt akft)
tty {cxt} t = tty? (AK⊎FT cxt) t

-- Now tcxt-ak is defined by
tcxt-ak .ε AKε =  (ε ,, U) --  inserting X:U at the beginning, and
tcxt-ak .(cxt ,, k) (AK,,k {cxt} ak k) =
  -- indexing all new kinds over x:X.
  (tcxt-ak cxt ak ,, Π (tnewty-ak cxt ak) (tkind? (inj₁ ak) k))

-- the type of the new variable x is just X, appropriately substituted.
tnewty-ak .ε AKε = (el (coe (U[] (kpop U)) ktop))
tnewty-ak .(cxt ,, k) (AK,,k {cxt} ak k) = tnewty-ak cxt ak [ kpop (Π (tnewty-ak cxt ak) (tkind? (inj₁ ak) k)) ]

-- tcxt-ft is defined by
tcxt-ft .(cxt , t) (FT,t₁ {cxt} ak t) =
  -- inserting x:X just before the first type
  ((tcxt-ak cxt ak , tnewty-ak cxt ak) , tty? (inj₁ ak) t)
-- and otherwise doing nothing to speak of
tcxt-ft .(cxt , t) (FT,t₊ {cxt} ft t) = (tcxt-ft cxt ft , tty? (inj₂ ft) t)
tcxt-ft .(cxt ,, k) (FT,,k {cxt} ft k) = (tcxt-ft cxt ft ,, tkind? (inj₂ ft) k)

-- A substitution from the extended context back to the original one,
-- which forgets X:U and x:X, while 'ap'plying all newly indexed
-- kind-variables to x to obtain values of their original kinds.
tcxt-sub? : (cxt : Cxt) → (akft : AllKinds cxt ⊎ FirstTy cxt) → Sub (tcxt? cxt akft) cxt

-- Now tkind? and tty? are easy to define by substitution.
tkind? {cxt} akft k = k [ tcxt-sub? cxt akft ]
tty? {cxt} akft t = t [ tcxt-sub? cxt akft ]

-- Finally we define this substitution.
-- When the original context is empty, we simply forget X:U and x:X.
tcxt-sub? .ε (inj₁ AKε) = kpop U ∘ tpop (el (coe (U[] (kpop U)) ktop))
-- The nontrivial case is when we add a new kind in between X and x.
tcxt-sub? .(cxt ,, k) (inj₁ (AK,,k {cxt} ak k)) =
  -- Here is a substitution which forgets X and x and *also* the newly added kind variable.
  (tcxt-sub? cxt (inj₁ ak) ∘ (kpop (Π (tnewty-ak cxt ak) (tkind? (inj₁ ak) k)) <id (tnewty-ak cxt ak)))
  -- Now we extend it by a term living in the original kind,
  -- obtained by 'ap'plying as above.  The hard part is proving that
  -- this application is valid; it requires knowing that Πs are
  -- preserved under substitution, and also that [] is functorial.
  ≪ coe (k[∘] _ _ _)
        (ap (coe (Π[] (tnewty-ak cxt ak) (tkind? (inj₁ ak) k) (kpop (Π (tnewty-ak cxt ak) (tkind? (inj₁ ak) k))))
                 ktop))
                     -- here is the variable which we 'ap'ply.
-- The other three cases are easy.
tcxt-sub? .(cxt , t) (inj₂ (FT,t₁ {cxt} ak t)) = tcxt-sub? cxt (inj₁ ak) <id t
tcxt-sub? .(cxt , t) (inj₂ (FT,t₊ {cxt} ft t)) = tcxt-sub? cxt (inj₂ ft) <id t
tcxt-sub? .(cxt ,, k) (inj₂ (FT,,k {cxt} ft k)) = tcxt-sub? cxt (inj₂ ft) ≪id k

-- No need for tsub, t≣, or ttm !?!

-- tcxt-ak always produces a contexts consisting only of kinds.
tcxt-ak-ak : (cxt : Cxt) → (ak : AllKinds cxt) → AllKinds (tcxt-ak cxt ak)
tcxt-ak-ak .ε AKε =  AK,,k AKε U
tcxt-ak-ak .(cxt ,, k) (AK,,k {cxt} ak k) =  AK,,k (tcxt-ak-ak cxt ak) (Π (tnewty-ak cxt ak) (k [ tcxt-sub? cxt (inj₁ ak) ]))

----------
-- Towers
----------

-- By iterating tcxt (or more precisely tcxt-ak), we produce contexts
-- that represent towers.
tower : ℕ → Cxt

-- These consist only of kinds (we need that to complete the iteration).
tower-ak : (n : ℕ) → AllKinds (tower n)

tower  = ε
tower (suc n) = tcxt-ak (tower n) (tower-ak n)

tower-ak  = AKε
tower-ak (suc n) = tcxt-ak-ak (tower n) (tower-ak n)

-- On my computer...
--   tower 3 terminates instantly.
--   tower 4 takes ~3 seconds.
--   tower 5 crashes Agda.

-- But 'extract (tower 10)' terminates instantly.
-- Even 'extract (tower 20)' takes only ~2 seconds.