module lec3 where

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

infixl  6 _`,_
infix   5 _⊢_ _⊑_ _⊢Ne_ _⊢Nf_ _⊢Env_
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[ _ ]

⊑-trans : Γ ⊑ Γ' → Γ' ⊑ Γ'' → Γ ⊑ Γ''
⊑-trans i        base      = i
⊑-trans i        (drop i') = drop (⊑-trans i i')
⊑-trans (drop i) (keep i') = ⊑-trans i (drop i')
⊑-trans (keep i) (keep i') = keep (⊑-trans i i')

-- 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

  lam   : Γ `, A ⊢ B
          ----------
        → Γ ⊢ A ⇒ B

  app   : Γ ⊢ A ⇒ B → Γ ⊢ A
          ------------
        → Γ ⊢ B

-- 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₂)
wkTm i (lam t)        = lam (wkTm (keep i) t)
wkTm i (app t u)      = app (wkTm i t) (wkTm i u)

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

data _⊢Ne_ : Ctx → Ty → Set
data _⊢Nf_ : Ctx → Ty → Set

-- Neutral terms
data _⊢Ne_ where
  var  : Var Γ A → Γ ⊢Ne A
  fst  : Γ ⊢Ne (A × B) → Γ ⊢Ne A
  snd  : Γ ⊢Ne (A × B) → Γ ⊢Ne B
  app  : Γ ⊢Ne (A ⇒ B) → Γ ⊢Nf A → Γ ⊢Ne B

-- Normal forms
data _⊢Nf_ 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
  lam   : (Γ `, A ⊢Nf B) → Γ ⊢Nf (A ⇒ B)

-- neutrals admit weakening
wkNe : Γ ⊑ Γ' → Γ ⊢Ne A → Γ' ⊢Ne A
-- normal forms admit weakening
wkNf : Γ ⊑ Γ' → Γ ⊢Nf A → Γ' ⊢Nf 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)
wkNe i (app n m) = app (wkNe i n) (wkNf i m)

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)
wkNf i (lam n)        = lam (wkNf (keep i) n)

-- 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)))
emb (A ⇒ B) n = lam (emb B (app (wkNe fresh n) (emb A (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
⊥' = λ Γ → ⊥

-- exponential family
_→'_ : Fam → Fam → Fam
X →' Y = λ Δ → ∀ {Δ'} → Δ ⊑ Δ' → X Δ' → Y Δ'

app' : (X →' Y) ×' X →̇ Y
app' (f , x) = f ⊑-refl[ _ ] x

app′ : X ×' (X →' Y) →̇ Y
app′ (x , f) = f ⊑-refl[ _ ] x

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

-- We can interpret the calculus
-- in families of sets given the
-- following abstract parmaeters
module Model
  -- a map of families
  (𝒯     : Fam → Fam)
  -- admits weakening
  (wk𝒯   : {X : Fam} → Weakens X → Weakens (𝒯 X))
  -- "monadic" operations
  (ret    : {X : Fam} → X →̇ 𝒯 X)
  (fmap   : {X Y : Fam} → (X →̇ Y) → 𝒯 X →̇ 𝒯 Y)
  (join   : {X : Fam} → 𝒯 (𝒯 X) →̇ 𝒯 X)
  (str[_] : {X : Fam} (wkX : Weakens X) → {Y : Fam} → X ×' 𝒯 Y →̇ 𝒯 (X ×' Y))
  -- a "valuation" for base types
  (V𝕓    : Fam)
  -- a "localising" operation
  (wkV𝕓  : Weakens V𝕓)
  (locV𝕓 : 𝒯 V𝕓 →̇ V𝕓)
  where

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

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

  wk⟦_⟧ : ∀ A → Weakens ⟦ A ⟧
  wk⟦ 𝕓 ⟧      i v = wkV𝕓 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⟦ A ⇒ B ⟧  i v = λ i' a → v (⊑-trans i i') a
  
  wk⟦_⟧ᶜ : ∀ Γ → Weakens ⟦ Γ ⟧ᶜ
  wk⟦ [] ⟧ᶜ       i γ       = tt
  wk⟦ (Γ `, a) ⟧ᶜ i (γ , v) = (wk⟦_⟧ᶜ Γ i γ) , wk⟦_⟧ a i v

  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
  loc (A ⇒ B) m = λ i a → loc B (fmap app′
    (str[_] wk⟦ A ⟧ (a , wk𝒯 wk⟦ A ⇒ B ⟧ i m)))

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

  module Interp 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[ wk⟦ Γ ⟧ᶜ ] (γ , vˢ)
      --
      match : ⟦ Γ ⟧ᶜ ×' (⟦ A ⟧ ⊎' ⟦ B ⟧) →̇ ⟦ C ⟧
      match (γ' , (inj₁ v₁)) = ⟦ t₁ ⟧ᵗ (γ' , v₁)
      match (γ' , (inj₂ v₂)) = ⟦ t₂ ⟧ᵗ (γ' , v₂)
    ⟦ lam {Γ} t ⟧ᵗ γ = λ i a → ⟦ t ⟧ᵗ (wk⟦ Γ ⟧ᶜ i γ , a)
    ⟦ app t u   ⟧ᵗ γ = app' (⟦ t ⟧ᵗ γ , ⟦ u ⟧ᵗ γ)

--
-- "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₂)

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 𝒞 wk𝒞 ret fmap join str (Nf' 𝕓) wkNf locNf -- gives ⟦_⟧, ⟦_⟧ᶜ, loc

open Interp

reflect : ∀ A → Ne' A →̇ ⟦ A ⟧
reify : ∀ A → ⟦ A ⟧ →̇ Nf' 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))))
reflect (A ⇒ B) n = λ i a → reflect B (app (wkNe i n) (reify A 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)
reify (A ⇒ B) f = lam (reify B (f fresh (reflect A (var zero))))

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[ Γ ]))

--------------------------------
-- Cherry-picked toy examples --
--------------------------------

pattern x₀ = var zero

↑₁ : Γ ⊢ A → (Γ `, B) ⊢ A
↑₁ = wkTm fresh

--
-- Term macros
--

Bool : Ty
Bool = 𝟙 + 𝟙

true false : Γ ⊢ Bool
true  = inl unit
false = inr unit

ifte : Γ ⊢ Bool → Γ ⊢ A → Γ ⊢ A → Γ ⊢ A
ifte s t u = case s (↑₁ t) (↑₁ u)

not : Γ ⊢ Bool → Γ ⊢ Bool
not s = ifte s false true

--
-- Identifying equal terms
--

id : ([] `, Bool) ⊢ Bool
id = x₀

‵not : ([] `, Bool) ⊢ Bool
‵not = not x₀

not-not : ([] `, Bool) ⊢ Bool
not-not = not (not x₀)

not-not-not : ([] `, Bool) ⊢ Bool
not-not-not = not not-not

--
-- Refuting existence of terms
--

no-ne : Impossible ([] ⊢Ne A)
no-ne (fst t)   = no-ne t
no-ne (snd t)   = no-ne t
no-ne (app t _) = no-ne t

consistent : Impossible ([] ⊢ 𝟘)
consistent t = no-nf-𝟘 (norm t)
  where
  no-nf-𝟘 : Impossible ([] ⊢Nf 𝟘)
  no-nf-𝟘 (abort x)    = no-ne x
  no-nf-𝟘 (case x _ _) = no-ne x