module CC.Term

import CC.Name
import CC.Thinning

import Control.Function

import Data.Nat
import Data.Singleton
import Data.SnocList.Elem

-- Definition ------------------------------------------------------------------

public export
Context : Type
Context = SnocList Name

public export
data IsVar : Nat -> Name -> Context -> Type where
  Here : (eq : m ~~ n) -> IsVar 0 m (ctx :< n)
  There : IsVar k n ctx -> IsVar (S k) n (ctx :< _)

public export
data Term : Context -> Type where
  Var : (k : Nat) -> (0 n : Name) -> (0 prf : IsVar k n ctx) -> Term ctx

  Let :
    (n : Name) ->
    (ty : Maybe (Term ctx)) ->
    (val : Term ctx) ->
    (t : Term (ctx :< n)) ->
    Term ctx

  Set : Term ctx

  Pi : (n : Maybe Name) -> Term ctx -> Term (ctx ++ maybe [<] ([<] :<) n) -> Term ctx
  Abs : (n : Name) -> Term (ctx :< n) -> Term ctx
  App : Term ctx -> Term ctx -> Term ctx

public export
data Value : Context -> Type
public export
data Neutral : Context -> Type

public export
data Env : (local, global : Context) -> Type where
  Lin  : Env [<] global
  (:<) : Env local global -> (p : (Name, Lazy (Value global))) -> Env (local :< fst p) global

data Value where
  Ntrl : Neutral ctx -> Value ctx
  VSet : Value ctx
  VPi :
    (n : Name) ->
    Lazy (Value ctx) ->
    Env local ctx ->
    Term (local :< n) ->
    Value ctx
  VFunc : Env local ctx -> Term local -> Term local -> Value ctx
  VAbs : Env local ctx -> (n : Name) -> Term (local :< n) -> Value ctx

data Neutral where
  VVar : (k : Nat) -> (0 n : Name) -> (0 prf : IsVar k n ctx) -> Neutral ctx
  VApp : Neutral ctx -> Lazy (Value ctx) -> Neutral ctx

%name IsVar prf
%name Term t, u
%name Value v
%name Neutral n
%name Env env

-- Operations ------------------------------------------------------------------

export
index : Env local global -> (k : Nat) -> (0 prf : IsVar k n local) -> Value global
index (env :< (n, v)) 0 (Here eq) = v
index (env :< _) (S k) (There prf) = index env k prf

export
fromElem : {n : Name} -> (prf : Elem n ctx) -> IsVar (elemToNat prf) n ctx
fromElem Here = Here (reflexive {rel = (~~)})
fromElem (There i) = There (fromElem i)

export
tabulate :
  {local : Context} ->
  (f : Nat -> Nat) ->
  (0 prf : {k, n : _} -> IsVar k n local -> IsVar (f k) n global) ->
  Env local global
tabulate {local = [<]} f prf = [<]
tabulate {local = _ :< n} f prf =
  tabulate (f . S) (prf . There) :<
  (n, Ntrl $ VVar (f 0) n (prf $ fromElem Here))

export
(++) : Env local1 global -> Env local2 global -> Env (local1 ++ local2) global
env ++ [<] = env
env ++ env' :< p = (env ++ env') :< p

recomputeLocal : Env local global -> Singleton local
recomputeLocal [<] = Val [<]
recomputeLocal (env :< (n, _)) = Val (:< n) <*> recomputeLocal env

-- Weakening -------------------------------------------------------------------

wknVar' :
  forall src, tgt.
  {0 thin : src `Thins` tgt} ->
  (k : Nat) ->
  (0 prf : IsVar k n src) ->
  View thin ->
  Nat
wknVar' k prf (Drop thin) = S (wknVar' k prf (view thin))
wknVar' 0 (Here eq) (Keep thin) = 0
wknVar' (S k) (There prf) (Keep thin) = S (wknVar' k prf (view thin))

wknIsVar' :
  forall src, tgt, k.
  {0 thin : src `Thins` tgt} ->
  (prf : IsVar k n src) ->
  (view : View thin) ->
  IsVar (wknVar' k prf view) n tgt
wknIsVar' prf (Drop thin) = There (wknIsVar' prf (view thin))
wknIsVar' (Here eq) (Keep thin) = Here eq
wknIsVar' (There prf) (Keep thin) = There (wknIsVar' prf (view thin))

export
wknVar : src `Thins` tgt -> {k : Nat} -> (0 prf : IsVar k n src) -> Nat
wknVar thin {k} prf = wknVar' k prf (view thin)

export
wknIsVar :
  (prf : IsVar k n src) ->
  (thin : src `Thins` tgt) ->
  IsVar (wknVar thin prf) n tgt
wknIsVar prf thin = wknIsVar' prf (view thin)

export
wknTerm : Term src -> src `Thins` tgt -> Term tgt
wknTerm (Var k n prf) thin = Var (wknVar thin prf) n (wknIsVar prf thin)
wknTerm (Let n Nothing val t) thin = Let n Nothing (wknTerm val thin) (wknTerm t (keep thin))
wknTerm (Let n (Just ty) val t) thin =
  Let n (Just $ wknTerm ty thin) (wknTerm val thin) (wknTerm t (keep thin))
wknTerm Set thin = Set
wknTerm (Pi Nothing dom cod) thin = Pi Nothing (wknTerm dom thin) (wknTerm cod thin)
wknTerm (Pi (Just n) dom cod) thin = Pi (Just n) (wknTerm dom thin) (wknTerm cod (keep thin))
wknTerm (Abs n t) thin = Abs n (wknTerm t (keep thin))
wknTerm (App t u) thin = App (wknTerm t thin) (wknTerm u thin)

export
wknNtrl : Neutral src -> src `Thins` tgt -> Neutral tgt
export
wknVal : Value src -> src `Thins` tgt -> Value tgt
export
wknEnv : Env local src -> src `Thins` tgt -> Env local tgt

wknNtrl (VVar k n prf) thin = VVar (wknVar thin prf) n (wknIsVar prf thin)
wknNtrl (VApp n v) thin = VApp (wknNtrl n thin) (wknVal v thin)

wknVal (Ntrl n) thin = Ntrl (wknNtrl n thin)
wknVal VSet thin = VSet
wknVal (VPi n dom env cod) thin = VPi n (wknVal dom thin) (wknEnv env thin) cod
wknVal (VFunc env dom cod) thin = VFunc (wknEnv env thin) dom cod
wknVal (VAbs env n t) thin = VAbs (wknEnv env thin) n t

wknEnv [<] thin = [<]
wknEnv (env :< (n, v)) thin = wknEnv env thin :< (n, wknVal v thin)

-- Renaming --------------------------------------------------------------------

renVar : {src, tgt : Context} -> length src = length tgt -> IsVar k n src -> (m ** IsVar k m tgt)
renVar prf (Here eq) {tgt = _ :< m} = (m ** fromElem Here)
renVar prf (There prf') {tgt = _ :< m} =
  let rec = renVar (injective prf) prf' in (rec.fst ** There rec.snd)

export
renTerm : length src = length tgt -> Term src -> Term tgt
renTerm prf (Var k n prf') = let 0 var = renVar prf prf' in Var k var.fst var.snd
renTerm prf (Let n Nothing val t) =
  Let n Nothing (renTerm prf val) (renTerm (cong S prf) t)
renTerm prf (Let n (Just ty) val t) =
  Let n (Just $ renTerm prf ty) (renTerm prf val) (renTerm (cong S prf) t)
renTerm prf Set = Set
renTerm prf (Pi Nothing t u) = Pi Nothing (renTerm prf t) (renTerm prf u)
renTerm prf (Pi (Just n) t u) = Pi (Just n) (renTerm prf t) (renTerm (cong S prf) u)
renTerm prf (Abs n t) = Abs n (renTerm (cong S prf) t)
renTerm prf (App t u) = App (renTerm prf t) (renTerm prf u)

-- Lifting ---------------------------------------------------------------------

inrVar : IsVar k n ctx -> IsVar k n (global ++ ctx)
inrVar (Here eq) = Here eq
inrVar (There prf) = There (inrVar prf)

liftEnv :
  {global : Context} ->
  (ctx : Context) ->
  Env local global ->
  Env (local ++ ctx) (global ++ ctx)
liftEnv [<] env = env
liftEnv ctx env = wknEnv env (id ++ empty) ++ tabulate id inrVar