module WellTypedTerms where

-- Utilities
data _≅_ {A : Set} : {A' : Set} → A → A' → Set where
 refl : {a : A} → a ≅ a

trans : ∀{A A' A''}{a : A}{a' : A'}{a'' : A''} → a ≅ a' → a' ≅ a'' → a ≅ a''
trans refl p = p

sym : ∀{A A'}{a : A}{a' : A'} → a ≅ a' → a' ≅ a
sym refl = refl

resp : ∀{A}{B : A → Set}(f : ∀ a → B a){a a' : A} → a ≅ a' → f a ≅ f a'
resp f refl = refl

resp2 : ∀{A}{B : A → Set}{C : (x : A) → B x → Set}(f : (x : A)(y : B x) → C x y)
          {a a' : A} → a ≅ a' → {b : B a}{b' : B a'} → b ≅ b' → f a b ≅ f a' b'
resp2 f refl refl = refl

_•_ : {A : Set}{B : A → Set}{C : (a : A) → B a → Set} →
     (∀ {a} b → C a b) → (g : (∀ (a : A) → B a)) → ∀ (a : A) → C a (g a)
f • g = λ a → f (g a)

id : {A : Set} → A → A
id a = a

data ⊥ : Set where

⊥-elim : ∀ {A : Set} -> ⊥ -> A
⊥-elim ()
record ⊤ : Set where
  constructor tt
data _+_ (S T : Set) : Set where
  inl : S -> S + T
  inr : T -> S + T

--Well type lambda terms

data Ty : Set where
 ι   : Ty
 _⇒_ : Ty → Ty → Ty

data Con : Set where
 ε   : Con
 _<_ : Con → Ty → Con

_<<_ : Con -> Con -> Con
Γ << ε = Γ
Γ << (Δ < τ) = (Γ << Δ) < τ

data Var : Con → Ty → Set where
 vz : ∀{Γ σ} → Var (Γ < σ) σ
 vs : ∀{Γ σ τ} → Var Γ σ → Var (Γ < τ) σ

data Tm : Con → Ty → Set where
 var : ∀{Γ σ} → Var Γ σ → Tm Γ σ
 app : ∀{Γ σ τ} → Tm Γ (σ ⇒ τ) → Tm Γ σ → Tm Γ τ
 lam : ∀{Γ σ τ} → Tm (Γ < σ) τ → Tm Γ (σ ⇒ τ)

-- Substitutions

Sub : Con → Con → Set
Sub Γ Δ = ∀ K {σ} → Var (Γ << K) σ → Tm (Δ << K) σ

sub : ∀{Γ Δ} → Sub Γ Δ → ∀ K {σ} → Tm (Γ << K) σ → Tm (Δ << K) σ
sub f K (var x)   = f K x
sub f K (app t u) = app (sub f K t) (sub f K u)
sub f K (lam t)   = lam (sub f (K < _) t)

subId : ∀{Γ} → Sub Γ Γ
subId = λ K x -> var x

subComp : ∀{B Γ Δ} → Sub Γ Δ → Sub B Γ → Sub B Δ
subComp f g K = sub f K • g K

-- first monad law, the second holds definitionally
subid : ∀{Γ}K{σ}(t : Tm (Γ << K) σ) → sub subId K t ≅ id t
subid K (var x)   = refl
subid K (app t u) = resp2 app (subid K t) (subid K u)
subid K (lam t)   = resp lam (subid (K < _) t)

-- third monad law
subcomp : ∀{B Γ Δ}(f : Sub Γ Δ)(g : Sub B Γ)K{σ}(t : Tm (B << K) σ) →
         sub (subComp f g) K t ≅ (sub f K • sub g K) t
subcomp f g K (var x)   = refl
subcomp f g K (app t u) = resp2 app (subcomp f g K t) (subcomp f g K u)
subcomp f g K (lam t)   = resp lam (subcomp f g (K < _) t)

-- Renaming

Ren : Con → Con → Set
Ren Γ Δ = ∀ K {σ} → Var (Γ << K) σ → Var (Δ << K) σ

renId : ∀{Γ} → Ren Γ Γ
renId = \ _ x -> x

renComp : ∀{B Γ Δ} → Ren Γ Δ → Ren B Γ → Ren B Δ
renComp f g = \ K -> f K • g K

ren : ∀{Γ Δ} → Ren Γ Δ → ∀ K {σ} → Tm (Γ << K) σ → Tm (Δ << K) σ
ren f K x = sub (λ K → subId K • f K) K x

renid : ∀{Γ} K {σ}(t : Tm (Γ << K) σ) → ren renId K t ≅ id t
renid = subid

rencomp : ∀ {B Γ Δ}(f : Ren Γ Δ)(g : Ren B Γ) K {σ}(t : Tm (B << K) σ) →
           ren (renComp f g) K t ≅ (ren f K • ren g K) t
rencomp f g K (var x)   = refl
rencomp f g K (app t u) = resp2 app (rencomp f g K t) (rencomp f g K u)
rencomp {B}{Γ}{Δ} f g K (lam t) = resp lam (rencomp f g (K < _) t)

-- these are now corollaries
rensub : ∀{B Γ Δ}(f : Ren Γ Δ)(g : Sub B Γ)K{σ}(t : Tm (B << K) σ) →
        (ren f K • sub g K) t ≅ sub (λ K -> ren f K • g K) K t
rensub f g K t = sym (subcomp (λ K -> subId K • f K) g K t)

subren : ∀{B Γ Δ}(f : Sub Γ Δ)(g : Ren B Γ)K{σ}(t : Tm (B << K) σ) →
        (sub f K • ren g K) t ≅ sub (λ K -> f K • g K) K t
subren f g K x = sym (subcomp f (λ K → subId K • g K) K x)

data Case (τ : _) (Γ Δ : Con) : Set where
  inj₁ : (∀ K -> Var (K << Γ) τ) -> Case τ Γ Δ
  inj₂ : Var Δ τ -> Case τ Γ Δ

case : ∀ {τ} Γ {Δ} -> Var (Δ << Γ) τ -> Case τ Γ Δ
case ε x = inj₂ x
case {τ} (Γ < .τ) vz = inj₁ (λ K → vz)
case (Γ < σ) (vs x) with case Γ x 
case (Γ < σ) (vs x) | inj₁ τ∈Γ = inj₁ λ K → vs (τ∈Γ K)
case (Γ < σ) (vs x) | inj₂ τ∈Δ = inj₂ τ∈Δ

_≪_ : ∀ {Δ τ} -> Var Δ τ -> ∀ Γ -> Var (Δ << Γ) τ
x ≪ ε = x
x ≪ (Γ < σ) = vs (x ≪ Γ)

Con-Ren : ∀ Γ {Δ} -> Ren Δ (Δ << Γ)
Con-Ren Γ K x with case K x
... | inj₁ τ∈K = τ∈K _
... | inj₂ τ∈Δ = (τ∈Δ ≪ Γ) ≪ K

lift : ∀ {Γ Δ} -> (∀ τ -> Var Γ τ -> Tm Δ τ) -> Sub Γ Δ
lift f K x with case K x
lift f K x | inj₁ τ∈K = var (τ∈K _)
lift f K x | inj₂ τ∈Γ = ren (Con-Ren K) ε (f _ τ∈Γ)

-- NbE inspired by http://personal.cis.strath.ac.uk/~conor/fooling/Nobby2.agda but typed
mutual
  data Nm : Con -> Ty -> Set where
    !_ : ∀ {Γ τ} -> Ne Γ τ -> Nm Γ τ 
    lam : ∀{Γ σ τ} → Nm (Γ < σ) τ → Nm Γ (σ ⇒ τ)

  data Ne : Con -> Ty -> Set where
    var : ∀{Γ σ} → Var Γ σ → Ne Γ σ
    app : ∀{Γ σ τ} → Ne Γ (σ ⇒ τ) → Nm Γ σ → Ne Γ τ

mutual
  [_] : Ty -> Con -> Set
  [ T ] Γ = (∀ {Δ} -> Ren Γ Δ -> Ne Δ T) + (∀ {Δ} -> Ren Γ Δ -> [ T ]act Δ)

  [_]act : Ty -> Con -> Set
  [ ι ]act  Γ = ⊥
  [ S ⇒ T ]act Γ = [ S ] Γ -> [ T ] Γ

ren[] : ∀ {Γ Δ} -> Ren Γ Δ -> ∀ {τ} -> [ τ ] Γ -> [ τ ] Δ
ren[] le (inl y) = inl (λ le' -> y (renComp le' le))
ren[] le (inr y) = inr (λ le' -> y (renComp le' le))

mutual
  quo : ∀ {Γ} (τ : Ty) -> [ τ ] Γ -> Nm Γ τ
  quo ι (inr x) = ⊥-elim (x renId)
  quo ι (inl e) = ! (e renId)
  quo (S ⇒ T) f = lam (quo T (ren[] (Con-Ren (ε < _)) f $$ (inl (λ le -> var (le ε vz)))))
  
  _$$_ : forall {Γ S T} -> [ S ⇒ T ] Γ -> [ S ] Γ -> [ T ] Γ
  _$$_ {S = S} (inl fn) s = inl λ le -> app (fn le) (quo S (ren[] le s))
  _$$_ (inr fa) s = fa renId s

data Env : Con -> Con -> Set where
  ε : ∀ {Γ} -> Env ε Γ 
  _<_ : ∀ {G Γ S} -> Env G Γ -> [ S ] Γ -> Env (G < S) Γ

renEnv : ∀ {G Γ Δ} -> Ren Γ Δ -> Env G Γ -> Env G Δ
renEnv le ε = ε
renEnv le (g < s) = renEnv le g < ren[] le s

get : forall {G Γ T} -> Var G T -> Env G Γ -> [ T ] Γ
get vz     (g < s) = s
get (vs x) (g < t) = get x g

ev : forall {G Γ T} -> Tm G T -> Env G Γ -> [ T ] Γ
ev (var x) g = get x g
ev (app f s) g = ev f g $$ ev s g
ev (lam t) g = inr λ le s → ev t (renEnv le g < s)

nm : forall {G Γ τ} -> Tm G τ -> Env G Γ -> Nm Γ τ 
nm t g = quo _ (ev t g)