module lec2 where

open import Function using (_∘_)
open import Data.Unit  using (⊤ ; tt)
open import Data.Empty  using (⊥ ; ⊥-elim)
open import Data.Product using (Σ ; _,_ ; proj₁ ; proj₂ ; uncurry) renaming (_×_  to _‵×_)
open import Data.Sum using (_⊎_ ; inj₁ ; inj₂ ; [_,_])

infixl  6 _`,_
infix   5 _⊑_
infixr 10 _→̇_

-----------
-- Types --
-----------

data Ty : Set where
  𝕓 𝟙 𝟘   : Ty
  _×_ _+_ : Ty → Ty → Ty

variable
    A B C D : Ty

---------------------
-- Typing contexts --
---------------------

data Ctx : Set where
  []   : Ctx
  _`,_ : Ctx → Ty → Ctx

variable
  Γ Γ' Δ Δ' : Ctx

-----------------------------
-- Variables and weakening --
-----------------------------

-- De Bruijn indices
data Var : Ctx → Ty → Set where
  zero : Var (Γ `, A) A
  succ : (v : Var Γ A) → Var (Γ `, B) A

-- Weakening relation
data _⊑_  : Ctx → Ctx → Set where
  base : [] ⊑ []
  drop : (w : Γ ⊑ Δ) → Γ ⊑ Δ `, A
  keep : (w : Γ ⊑ Δ) → Γ `, A ⊑ Δ `, A

-- the relation ⊑ is reflexive
⊑-refl[_] : (Γ : Ctx) → Γ ⊑ Γ
⊑-refl[ []      ] = base
⊑-refl[ Γ `, _a ] = keep ⊑-refl[ Γ ]

-- named such since it corresponds to
-- weakening with a "fresh" variable
fresh : Γ ⊑ Γ `, A
fresh = drop ⊑-refl[ _ ]

-- variables admit weakening
wkVar : Γ ⊑ Γ' → Var Γ A → Var Γ' A
wkVar (drop i) v        = succ (wkVar i v)
wkVar (keep i) zero     = zero
wkVar (keep i) (succ v) = succ (wkVar i v)

-----------
-- Terms --
-----------

data _⊢_ : Ctx → Ty → Set where

  var   : Var Γ A
          -------
        → Γ ⊢ A

  unit  : Γ ⊢ 𝟙

  abort : Γ ⊢ 𝟘
          ------
        → Γ ⊢ A

  fst   : Γ ⊢ (A × B)
          -----------
        → Γ ⊢ A

  snd   : Γ ⊢ (A × B)
          -----------
        → Γ ⊢ B

  pair  : Γ ⊢ A → Γ ⊢ B
          -------------
        → Γ ⊢ (A × B)

  inl   : Γ ⊢ A
          -----------
        → Γ ⊢ (A + B)

  inr   : Γ ⊢ B
          -----------
        → Γ ⊢ (A + B)

  case  : Γ ⊢ (A + B) → (Γ `, A) ⊢ C → (Γ `, B) ⊢ C
          -----------------------------------------
        → Γ ⊢ C

-- terms admit weakening
wkTm : Γ ⊑ Γ' → Γ ⊢ A → Γ' ⊢ A
wkTm i (var x)        = var (wkVar i x)
wkTm i unit           = unit
wkTm i (abort t)      = abort (wkTm i t)
wkTm i (fst t)        = fst (wkTm i t)
wkTm i (snd t)        = snd (wkTm i t)
wkTm i (pair t₁ t₂)   = pair (wkTm i t₁) (wkTm i t₂)
wkTm i (inl t)        = inl (wkTm i t)
wkTm i (inr t)        = inr (wkTm i t)
wkTm i (case s t₁ t₂) = case (wkTm i s) (wkTm (keep i) t₁) (wkTm (keep i) t₂)

------------------------------------
-- Neutral terms and normal forms --
------------------------------------

-- Neutral terms
data _⊢Ne_ : Ctx → Ty → Set where
  var  : Var Γ A → Γ ⊢Ne A
  fst  : Γ ⊢Ne (A × B) → Γ ⊢Ne A
  snd  : Γ ⊢Ne (A × B) → Γ ⊢Ne B

-- Normal forms
data _⊢Nf_ : Ctx → Ty → Set where
  up    : Γ ⊢Ne 𝕓 → Γ ⊢Nf 𝕓
  unit  : Γ ⊢Nf 𝟙
  abort : Γ ⊢Ne 𝟘 → Γ ⊢Nf A
  pair  : Γ ⊢Nf A → Γ ⊢Nf B → Γ ⊢Nf (A × B)
  inl   : Γ ⊢Nf A → Γ ⊢Nf (A + B)
  inr   : Γ ⊢Nf B → Γ ⊢Nf (A + B)
  case  : Γ ⊢Ne (A + B) → (Γ `, A) ⊢Nf C → (Γ `, B) ⊢Nf C → Γ ⊢Nf C

-- neutrals admit weakening
wkNe : Γ ⊑ Γ' → Γ ⊢Ne A → Γ' ⊢Ne A
wkNe i (var x)  = var (wkVar i x)
wkNe i (fst n)  = fst (wkNe i n)
wkNe i (snd n)  = snd (wkNe i n)

-- normal forms admit weakening
wkNf : Γ ⊑ Γ' → Γ ⊢Nf A → Γ' ⊢Nf A
wkNf i (up x)         = up (wkNe i x)
wkNf i unit           = unit
wkNf i (abort x)      = abort (wkNe i x)
wkNf i (pair n m)     = pair (wkNf i n) (wkNf i m)
wkNf i (inl n)        = inl (wkNf i n)
wkNf i (inr n)        = inr (wkNf i n)
wkNf i (case n m1 m2) = case (wkNe i n) (wkNf (keep i) m1) (wkNf (keep i) m2)

-- a neutral can be "η expanded" to a normal form
emb : ∀ A → Γ ⊢Ne A → Γ ⊢Nf A
emb 𝕓 n = up n
emb 𝟙 n = unit
emb 𝟘 n = abort n
emb (A × B) n = pair (emb A (fst n)) (emb B (snd n))
emb (A + B) n = case n (inl (emb A (var zero))) (inr (emb B (var zero)))

------------------
-- Environments --
------------------

-- an environment is a "snoclist" of neutrals
data _⊢Env_ : Ctx → Ctx → Set where
  nil  : Δ ⊢Env []
  cons : Δ ⊢Env Γ → Δ ⊢Ne A → Δ ⊢Env (Γ `, A)

-- Obs: `Env Δ Γ` is an environment
-- for Γ (whose neutrals are) typed using Δ

-- environments admit weakening
wkEnv : Δ ⊑ Δ' → Δ ⊢Env Γ → Δ' ⊢Env Γ
wkEnv i nil        = nil
wkEnv i (cons γ x) = cons (wkEnv i γ) (wkNe i x)

-- the relation "_⊢Env_" is reflexive
idEnv[_] : ∀ Δ → Δ ⊢Env Δ
idEnv[ [] ]     = nil
idEnv[ Δ `, x ] = cons (wkEnv fresh (idEnv[ Δ ])) (var zero)

----------------------
-- Families of sets --
----------------------

Fam : Set₁
Fam = Ctx → Set

-- I used a dot ( ̇) during the lectures,
-- but will use a prime (') here instead
Var' Tm' Ne' Nf' : Ty → Fam
Env' : Ctx → Fam

Var' a Γ = Var Γ a
Tm'  a Γ = Γ ⊢ a
Ne'  a Γ = Γ ⊢Ne a
Nf'  a Γ = Γ ⊢Nf a
Env' Δ Γ = Γ ⊢Env Δ

private
  variable
    X Y : Ctx → Set

-- family of functions
_→̇_ : Fam → Fam → Set
_→̇_ X Y = {Δ : Ctx} → X Δ → Y Δ

-- product family
_×'_ : Fam → Fam → Fam
_×'_ X Y Γ = X Γ ‵× Y Γ

-- sum family
_⊎'_ : Fam → Fam → Fam
_⊎'_ X Y Γ = X Γ ⊎ Y Γ

-- unit family
⊤' : Fam
⊤' = λ Γ → ⊤

-- empty family
⊥' : Fam
⊥' = λ Γ → ⊥

-- We can interpret the calculus
-- in families of sets given the
-- following abstract parmaeters
module Model
  -- a map of families
  -- (I used `𝒞` instead of `𝒯` in lectures)
  (𝒯     : Fam → Fam)
  -- "monadic" operations
  (ret   : {X : Fam} → X →̇ 𝒯 X)
  (fmap  : {X Y : Fam} → (X →̇ Y) → 𝒯 X →̇ 𝒯 Y)
  (join  : {X : Fam} → 𝒯 (𝒯 X) →̇ 𝒯 X)
  -- a "valuation" for base types
  (V𝕓    : Fam)
  -- a "localising" operation
  (locV𝕓 : 𝒯 V𝕓 →̇ V𝕓)
  where

  ⟦_⟧ : Ty → Fam
  ⟦ 𝕓 ⟧ = V𝕓
  ⟦ 𝟙 ⟧ = ⊤'
  ⟦ 𝟘 ⟧ = 𝒯 ⊥'
  ⟦ A × B ⟧ = ⟦ A ⟧ ×' ⟦ B ⟧
  ⟦ A + B ⟧ = 𝒯 (⟦ A ⟧ ⊎' ⟦ B ⟧)

  ⟦_⟧ᶜ : Ctx → Fam
  ⟦ [] ⟧ᶜ     Γ = ⊤
  ⟦ Δ `, a ⟧ᶜ Γ = ⟦ Δ ⟧ᶜ Γ ‵× ⟦ a ⟧ Γ

  loc : ∀ A → 𝒯 ⟦ A ⟧ →̇ ⟦ A ⟧
  loc 𝕓 m = locV𝕓 m
  loc 𝟙 m = tt
  loc 𝟘 m = join m
  loc (A × B) m = loc A (fmap proj₁ m) , loc B (fmap proj₂ m)
  loc (A + B) m = join m

  lookup : Var Γ A → (⟦ Γ ⟧ᶜ →̇ ⟦ A ⟧)
  lookup zero     (_ , x) = x
  lookup (succ x) (γ , _) = lookup x γ

  --
  -- Interpretation of terms demands a `str` operation
  --
  -- I will get rid of this seemingly adhoc requirement in the next
  -- lecture by interpreting types and contexts in families
  -- that admit weakening (defined further below as `Weakens`)
  --
  -- I didn't here already do that here simply to highlight
  -- that we strictly don't need to work with such families yet.
  --
  module Interp (str[_]  : (Γ : Ctx) {X : Fam} → ⟦ Γ ⟧ᶜ ×' 𝒯 X →̇ 𝒯 (⟦ Γ ⟧ᶜ ×' X)) where

    ⟦_⟧ᵗ : Γ ⊢ A → (⟦ Γ ⟧ᶜ →̇ ⟦ A ⟧)
    ⟦ var x ⟧ᵗ           γ = lookup x γ
    ⟦ unit ⟧ᵗ            γ = tt
    ⟦ abort {Γ} {A} t ⟧ᵗ γ = loc A (fmap ⊥-elim (⟦ t ⟧ᵗ γ))
    ⟦ fst t ⟧ᵗ           γ = proj₁ (⟦ t ⟧ᵗ γ)
    ⟦ snd t ⟧ᵗ           γ = proj₂ (⟦ t ⟧ᵗ γ)
    ⟦ pair t u ⟧ᵗ        γ = ⟦ t ⟧ᵗ γ , ⟦ u ⟧ᵗ γ
    ⟦ inl t ⟧ᵗ           γ = ret (inj₁ (⟦ t ⟧ᵗ γ))
    ⟦ inr t ⟧ᵗ           γ = ret (inj₂ (⟦ t ⟧ᵗ γ))
    ⟦ case {Γ} {A} {B} {C} s t₁ t₂ ⟧ᵗ {Δ} γ = loc C (fmap match gvˢ)
      where
      -- value of the scrutinee
      vˢ : 𝒯 (⟦ A ⟧ ⊎' ⟦ B ⟧) Δ
      vˢ = ⟦ s ⟧ᵗ γ
      -- paired with environment
      gvˢ : 𝒯 (⟦ Γ ⟧ᶜ ×' (⟦ A ⟧ ⊎' ⟦ B ⟧)) Δ
      gvˢ = str[ Γ ] (γ , vˢ)
      --
      match : ⟦ Γ ⟧ᶜ ×' (⟦ A ⟧ ⊎' ⟦ B ⟧) →̇ ⟦ C ⟧
      match (γ' , (inj₁ v₁)) = ⟦ t₁ ⟧ᵗ (γ' , v₁)
      match (γ' , (inj₂ v₂)) = ⟦ t₂ ⟧ᵗ (γ' , v₂)

--
-- "Covering family"
--
-- Intuitively 𝒞 is a tree whose nodes denote branching
-- caused by case splitting and leaves denote values of
-- some set or an abort pretending to be a value
--
data 𝒞 (X : Fam) : Fam where
  val   : X Γ → 𝒞 X Γ
  abort : Γ ⊢Ne 𝟘 → 𝒞 X Γ
  case  : Γ ⊢Ne (A + B) → 𝒞 X (Γ `, A) → 𝒞 X (Γ `, B) → 𝒞 X Γ

ret : X →̇ 𝒞 X
ret = val

fmap : (X →̇ Y) → 𝒞 X →̇ 𝒞 Y
fmap f (val a)        = val (f a)
fmap f (abort n)      = abort n
fmap f (case n m₁ m₂) = case n (fmap f m₁) (fmap f m₂)

join : 𝒞 (𝒞 X) →̇ 𝒞 X
join (val a)        = a
join (abort n)      = abort n
join (case n m₁ m₂) = case n (join m₁) (join m₂)

Weakens : Fam → Set
Weakens X = {Γ Γ' : Ctx} → Γ ⊑ Γ' → X Γ → X Γ'

module _ {X : Fam} (wkX : Weakens X) where

  -- if the leaves can be weakened
  wk𝒞 : Weakens (𝒞 X)
  wk𝒞 i (val x)        = val (wkX i x)
  wk𝒞 i (abort n)      = abort (wkNe i n)
  wk𝒞 i (case n m₁ m₂) = case (wkNe i n) (wk𝒞 (keep i) m₁) (wk𝒞 (keep i) m₂)

  -- "strength" demands `Weakens X`
  str : {Y : Fam} → X ×' 𝒞 Y →̇ 𝒞 (X ×' Y)
  str (x , m) = go x m
    where
    go : ∀ {Y Γ} → X Γ → 𝒞 Y Γ → 𝒞 (X ×' Y) Γ
    go a (val B)        = val (a , B)
    go a (abort n)      = abort n
    go a (case n m₁ m₂) = case n (go (wkX fresh a) m₁) (go (wkX fresh a) m₂)

-- fold a "tree" whose leaves contain
-- normal forms into a normal form
locNf : 𝒞 (Nf' A) →̇ Nf' A
locNf (val n)        = n
locNf (abort n)      = abort n
locNf (case n c₁ c₂) = case n (locNf c₁) (locNf c₂)

open Model 𝒞 ret fmap join (Nf' 𝕓) locNf -- gives ⟦_⟧, ⟦_⟧ᶜ, loc

wk⟦⟧ : ∀ A → Weakens ⟦ A ⟧
wk⟦⟧ 𝕓       i v = wkNf i v
wk⟦⟧ 𝟙       i v = tt
wk⟦⟧ 𝟘       i v = wk𝒞 (λ i' → ⊥-elim) i v
wk⟦⟧ (A × B) i v = wk⟦⟧ A i (proj₁ v) , wk⟦⟧ B i (proj₂ v)
wk⟦⟧ (A + B) i v = wk𝒞 (λ i' → [ inj₁ ∘ wk⟦⟧ A i' , inj₂ ∘ wk⟦⟧ B i' ]) i v

wk⟦⟧ᶜ : ∀ Γ → Weakens ⟦ Γ ⟧ᶜ
wk⟦⟧ᶜ []       i γ       = tt
wk⟦⟧ᶜ (Γ `, a) i (γ , v) = (wk⟦⟧ᶜ Γ i γ) , wk⟦⟧ a i v

open Interp (str ∘ wk⟦⟧ᶜ) -- gives ⟦_⟧ᵗ

reflect : ∀ A → Ne' A →̇ ⟦ A ⟧
reflect 𝕓 n = up n
reflect 𝟙 n = tt
reflect 𝟘 n = abort n
reflect (A × B) n = (reflect A (fst n) , reflect B (snd n))
reflect (A + B) n = case n
  (val (inj₁ (reflect A (var zero))))
  (val (inj₂ (reflect B (var zero))))

reify : ∀ A → ⟦ A ⟧ →̇ Nf' A
reify 𝕓 x = x
reify 𝟙 x = unit
reify 𝟘 m = locNf (fmap ⊥-elim m)
reify (A × B) p = pair (reify A (proj₁ p)) (reify B (proj₂ p))
reify (A + B) m = locNf (fmap [ inl ∘ reify A , inr ∘ reify B ] m)

reflectᵉ : Env' Γ →̇ ⟦ Γ ⟧ᶜ
reflectᵉ nil        = tt
reflectᵉ (cons e n) = (reflectᵉ e) , reflect _ n

norm : Γ ⊢ A → Γ ⊢Nf A
norm {Γ} {A} t = reify A (⟦ t ⟧ᵗ (reflectᵉ idEnv[ Γ ]))