data Tm = Var Int | Ann Tm Tm | Abs Tm | App Tm Tm | Pi Tm Tm | Fix Tm | Uni
data Clos = Clos Tm Env
data Dm = DVar Int | DAbs Clos | DNeutral Int [Dm] | DPi Dm Clos | DFix Clos | DUni
type Env = [Dm]

capp :: Clos -> Dm -> Dm
capp (Clos b e) t = eval (t : e) b

vapp :: Dm -> Dm -> Dm
vapp a b =
  case a of
    DAbs c -> capp c b
    DFix c -> vapp (capp c a) b
    DVar i -> DNeutral i [b]
    DNeutral i as -> DNeutral i (b : as)
    tm -> error $ "impossible vapp"

eval :: Env -> Tm -> Dm
eval env Uni = DUni
eval env (Var i) = env !! i
eval env (Ann t ty) = eval env t
eval env (Abs b) = DAbs (Clos b env)
eval env (Pi t b) = DPi (eval env t) (Clos b env)
eval env (Fix b) = DFix (Clos b env)
eval env (App a b) = vapp (eval env a) (eval env b)

quote :: Int -> Dm -> Tm
quote k DUni = Uni
quote k (DVar i) = Var (k - (i + 1))
quote k (DAbs (Clos b e)) =
  Abs (quote (k + 1) (eval (DVar k : e) b))
quote k (DPi t (Clos b e)) =
  Pi (quote k t) (quote (k + 1) (eval (DVar k : e) b))
quote k (DFix (Clos b e)) =
  Fix (quote (k + 1) (eval (DVar k : e) b))
quote k (DNeutral i as) =
  foldr (flip App) (Var (k - (i + 1))) (map (quote k) as)

nf :: Tm -> Tm
nf t = quote 0 (eval [] t)

eqtype :: Int -> Dm -> Dm -> Bool
eqtype k DUni DUni = True
eqtype k (DVar i) (DVar j) = i == j
eqtype k (DFix c) (DFix c') =
  let v = DVar k in
  eqtype (k + 1) (capp c v) (capp c' v)
eqtype k (DAbs c) (DAbs c') =
  let v = DVar k in
  eqtype (k + 1) (capp c v) (capp c' v)
eqtype k (DAbs c) t' =
  let v = DVar k in
  eqtype (k + 1) (capp c v) (vapp t' v)
eqtype k t' (DAbs c) =
  let v = DVar k in
  eqtype (k + 1) (vapp t' v) (capp c v)
eqtype k v@(DFix c) t' =
  eqtype (k + 1) (capp c v) t'
eqtype k t' v@(DFix c) =
  eqtype (k + 1) t' (capp c v)
eqtype k (DNeutral i as) (DNeutral i' as') =
  i == i' && and (zipWith (eqtype k) as as')
eqtype k (DPi t c) (DPi t' c') =
  let v = DVar k in
  eqtype k t t' && eqtype (k + 1) (capp c v) (capp c' v)
eqtype _ _ _ = False

type Err t = Either String t

err :: String -> Err t
err = Left

index :: Int -> [t] -> String -> Err t
index _ [] msg = err msg
index 0 (h : _) _ = return h
index i (_ : t) msg = index (i - 1) t msg

assert :: Bool -> String -> Err ()
assert b msg = if b then return () else err msg

check :: Env -> Env -> Int -> Tm -> Dm -> Err ()
check tenv venv k (Abs b) (DPi t b') =
  let v = DVar k in
  check (t : tenv) (v : venv) (k + 1) b (capp b' v)
check tenv venv k (Abs b) tm@(DFix b') =
  check tenv venv k (Abs b) (capp b' tm)
check tenv venv k tm@(Fix b) t =
  check (t : tenv) (eval venv tm : venv) (k + 1) b t
check tenv venv k t ty = do
  ty' <- synth tenv venv k t
  assert (eqtype k ty' ty) $ "type are not equal " ++ (show $ quote k ty') ++ " ~ " ++ (show $ quote k ty)

synthapp :: Env -> Env -> Int -> Tm -> Dm -> Tm -> Err Dm
synthapp tenv venv k tm ta b =
  case ta of
    DPi t c -> do
      check tenv venv k b t
      return $ capp c (eval venv b)
    DFix c -> synthapp tenv venv k tm (capp c ta) b
    ty -> err $ "not a pi type in " ++ (show tm) ++ ": " ++ (show $ quote k ty)

synth :: Env -> Env -> Int -> Tm -> Err Dm
synth tenv venv k Uni = return DUni
synth tenv venv k (Var i) = index i tenv $ "var " ++ (show i) ++ " not found"
synth tenv venv k (Ann t ty) = do
  check tenv venv k ty DUni
  let ty' = eval venv ty
  check tenv venv k t ty'
  return ty'
synth tenv venv k (Pi t b) = do
  check tenv venv k t DUni
  check (eval venv t : tenv) (DVar k : venv) (k + 1) b DUni
  return DUni
synth tenv venv k tm@(App a b) = do
  ta <- synth tenv venv k a
  synthapp tenv venv k tm ta b
synth tenv venv k t = err $ "cannot synth " ++ (show t)

-- pretty
instance Show Tm where
  show (Var i) = show i
  show (Ann t ty) = "(" ++ (show t) ++ " : " ++ (show ty) ++ ")"
  show (Abs b) = "(\\" ++ (show b) ++ ")"
  show (App a b) = "(" ++ (show a) ++ " " ++ (show b) ++ ")"
  show (Pi t b) = "(" ++ (show t) ++ " -> " ++ (show b) ++ ")"
  show (Fix b) = "(fix " ++ (show b) ++ ")"
  show Uni = "*"

-- testing
v :: Int -> Tm
v = Var

a :: Tm -> [Tm] -> Tm
a f as = foldl App f as

nat :: Tm -- fix r. (t:*) -> t -> (r -> t) -> t  :  *
nat = Ann (Fix (Pi Uni $ Pi (v 0) $ Pi (Pi (v 2) (v 2)) (v 2))) Uni

z :: Tm -- \t z s. z : nat
z = Ann (Abs $ Abs $ Abs $ v 1) nat

s :: Tm -- \n t z s. s n : nat
s = Ann (Abs $ Abs $ Abs $ Abs $ App (v 0) (v 3)) (Pi nat nat)

bool :: Tm -- (t:*) -> t -> t -> t
bool = Pi Uni $ Pi (v 0) $ Pi (v 1) (v 2)

tru :: Tm -- \t a b. a : bool
tru = Ann (Abs $ Abs $ Abs $ v 1) bool

fls :: Tm -- \t a b. b : bool
fls = Ann (Abs $ Abs $ Abs $ v 0) bool

bnot :: Tm -- \b. b bool fls tru : bool -> bool
bnot = Ann (Abs $ a (v 0) [bool, fls, tru]) (Pi bool bool)

neven :: Tm -- fix r. \n. n bool tru (\m. not (r m)) : nat -> bool
neven = Ann (Fix $ Abs $ a (v 0) [bool, tru, Abs (App bnot (App (v 2) (v 0)))]) (Pi nat bool)

term :: Tm -- even (s (s z)) ~> tru
term = App neven (App s (App s z))

main :: IO ()
main = do
  putStrLn $ show term
  putStrLn $ show $ fmap (quote 0) $ synth [] [] 0 term
  putStrLn $ show $ nf term