module LPTypechecker (Type, typecheck) where
import Parser (CTerm(..), ITerm(..), Name(..), Program)

import Control.Monad (liftM)
import Control.Monad.Error
import Control.Monad.State
import Data.Map (Map, fromList, toList)
import Text.PrettyPrint.HughesPJ hiding (parens)
import qualified Text.PrettyPrint.HughesPJ as PP

typecheck :: Program -> Either String Program
typecheck p = liftM (\(prog, _, _) -> prog) $ foldM checkTerm ([], lpve,lpte) p
    where checkTerm (prog, ve, te) (name, term) =
            case inferTerm (ve, te) term of
                Left err -> Left $ "typecheck failure in " ++ name ++ "\n" ++ err
                Right (annTerm, typ) -> Right ((name, annTerm) : prog,
                                               (Global name, evalTerm ve term) : ve,
                                               (Global name, typ) : te)
          -- Yes we're evaluating each term during typechecking, but laziness keeps
          -- things sane.
          evalTerm  env term = iEval term (env, [])
          inferTerm env term = iType 0 env term

type NameEnv v = [(Name, v)]
type Ctx inf = [(Name, inf)]
type State v inf = (Bool, String, NameEnv v, Ctx inf)

data Value
   =  VLam  (Value -> Value)
   |  VStar
   |  VPi Value (Value -> Value)
   |  VNeutral Neutral
   |  VNat
   |  VZero
   |  VSucc Value
   |  VNil Value
   |  VCons Value Value Value Value
   |  VVec Value Value
   |  VEq Value Value Value
   |  VRefl Value Value
   |  VFZero Value
   |  VFSucc Value Value
   |  VFin Value

data Neutral
   =  NFree  Name
   |  NApp  Neutral Value
   |  NNatElim Value Value Value Neutral
   |  NVecElim Value Value Value Value Value Neutral
   |  NEqElim Value Value Value Value Value Neutral
   |  NFinElim Value Value Value Value Neutral
  
type Result a = Either String a

type Env = [Value]

vapp :: Value -> Value -> Value
vapp (VLam f)      v  =  f v
vapp (VNeutral n)  v  =  VNeutral (NApp n v)

vfree :: Name -> Value
vfree n = VNeutral (NFree n)

cEval :: CTerm -> (NameEnv Value,Env) -> Value
cEval (Inf  ii)    d  =  iEval ii d
cEval (Lam  c)     d  =  VLam (\ x -> cEval c (((\(e, d) -> (e,  (x : d))) d)))

cEval Zero      d  = VZero
cEval (Succ k)  d  = VSucc (cEval k d)

cEval (Nil a)          d  =  VNil (cEval a d)
cEval (Cons a n x xs)  d  =  VCons  (cEval a d) (cEval n d)
                                       (cEval x d) (cEval xs d)

cEval (Refl a x)       d  =  VRefl (cEval a d) (cEval x d)

cEval (FZero n)    d  =  VFZero (cEval n d)
cEval (FSucc n f)  d  =  VFSucc (cEval n d) (cEval f d)

iEval :: ITerm -> (NameEnv Value,Env) -> Value
iEval (Ann  c _)       d  =  cEval c d

iEval Star           d  =  VStar   
iEval (Pi ty ty')    d  =  VPi (cEval ty d) (\ x -> cEval ty' (((\(e, d) -> (e,  (x : d))) d)))   

iEval (Free  x)      d  =  case lookup x (fst d) of Nothing ->  (vfree x); Just v -> v 
iEval (Bound  ii)    d  =  (snd d) !! ii
iEval (i :$: c)       d  =  vapp (iEval i d) (cEval c d)

iEval Nat                  d  =  VNat
iEval (NatElim m mz ms n)  d 
  =  let  mzVal = cEval mz d
          msVal = cEval ms d
          rec nVal =
            case nVal of
              VZero       ->  mzVal
              VSucc k     ->  msVal `vapp` k `vapp` rec k
              VNeutral n  ->  VNeutral
                               (NNatElim (cEval m d) mzVal msVal n)
              _            ->  error "internal: eval natElim"
     in   rec (cEval n d)

iEval (Vec a n)                 d  =  VVec (cEval a d) (cEval n d)

iEval (VecElim a m mn mc n xs)  d  =
  let  mnVal  =  cEval mn d
       mcVal  =  cEval mc d
       rec nVal xsVal =
         case xsVal of
           VNil _          ->  mnVal
           VCons _ k x xs  ->  foldl vapp mcVal [k, x, xs, rec k xs]
           VNeutral n      ->  VNeutral
                                (NVecElim  (cEval a d) (cEval m d)
                                            mnVal mcVal nVal n)
           _                ->  error "internal: eval vecElim"
  in   rec (cEval n d) (cEval xs d)

iEval (Eq a x y)                d  =  VEq (cEval a d) (cEval x d) (cEval y d)
iEval (EqElim a m mr x y eq)    d  =
  let  mrVal  =  cEval mr d
       rec eqVal =
         case eqVal of
           VRefl _ z -> mrVal `vapp` z
           VNeutral n ->  
             VNeutral (NEqElim  (cEval a d) (cEval m d) mrVal
                                  (cEval x d) (cEval y d) n)
           _ -> error "internal: eval eqElim"
  in   rec (cEval eq d)

iEval (Fin n)                d  =  VFin (cEval n d)
iEval (FinElim m mz ms n f)  d  =
  let  mzVal  =  cEval mz d
       msVal  =  cEval ms d
       rec fVal =
         case fVal of
           VFZero k        ->  mzVal `vapp` k
           VFSucc k g      ->  foldl vapp msVal [k, g, rec g]
           VNeutral n'     ->  VNeutral
                                (NFinElim  (cEval m d) (cEval mz d)
                                            (cEval ms d) (cEval n d) n')
           _                ->  error "internal: eval finElim"
  in   rec (cEval f d)

iSubst :: Int -> ITerm -> ITerm -> ITerm
iSubst ii i'   (Ann c c')     =  Ann (cSubst ii i' c) (cSubst ii i' c')


iSubst ii r  Star           =  Star  
iSubst ii r  (Pi ty ty')    =  Pi  (cSubst ii r ty) (cSubst (ii + 1) r ty')

iSubst ii i' (Bound j)      =  if ii == j then i' else Bound j
iSubst ii i' (Free y)       =  Free y
iSubst ii i' (i :$: c)       =  iSubst ii i' i :$: cSubst ii i' c

iSubst ii r  Nat            =  Nat
iSubst ii r  (NatElim m mz ms n)
                              =  NatElim (cSubst ii r m)
                                          (cSubst ii r mz) (cSubst ii r ms)
                                          (cSubst ii r ms)

iSubst ii r  (Vec a n)      =  Vec (cSubst ii r a) (cSubst ii r n)
iSubst ii r  (VecElim a m mn mc n xs)
                              =  VecElim (cSubst ii r a) (cSubst ii r m)
                                          (cSubst ii r mn) (cSubst ii r mc)
                                          (cSubst ii r n) (cSubst ii r xs)

iSubst ii r  (Eq a x y)     =  Eq (cSubst ii r a)
                                     (cSubst ii r x) (cSubst ii r y)
iSubst ii r  (EqElim a m mr x y eq)
                              =  VecElim (cSubst ii r a) (cSubst ii r m)
                                          (cSubst ii r mr) (cSubst ii r x)
                                          (cSubst ii r y) (cSubst ii r eq)

iSubst ii r  (Fin n)        =  Fin (cSubst ii r n)
iSubst ii r  (FinElim m mz ms n f)
                              =  FinElim (cSubst ii r m)
                                          (cSubst ii r mz) (cSubst ii r ms)
                                          (cSubst ii r n) (cSubst ii r f)

cSubst :: Int -> ITerm -> CTerm -> CTerm
cSubst ii i' (Inf i)      =  Inf (iSubst ii i' i)
cSubst ii i' (Lam c)      =  Lam (cSubst (ii + 1) i' c)

cSubst ii r  Zero         =  Zero
cSubst ii r  (Succ n)     =  Succ (cSubst ii r n)

cSubst ii r  (Nil a)      =  Nil (cSubst ii r a)
cSubst ii r  (Cons a n x xs)
                            =  Cons (cSubst ii r a) (cSubst ii r x)
                                     (cSubst ii r x) (cSubst ii r xs)

cSubst ii r  (Refl a x)   =  Refl (cSubst ii r a) (cSubst ii r x)

cSubst ii r  (FZero n)    =  FZero (cSubst ii r n)
cSubst ii r  (FSucc n k)  =  FSucc (cSubst ii r n) (cSubst ii r k)

quote :: Int -> Value -> CTerm
quote ii (VLam t)
  =     Lam (quote (ii + 1) (t (vfree (Quote ii))))


quote ii VStar = Inf Star 
quote ii (VPi v f)                                       
    =  Inf (Pi (quote ii v) (quote (ii + 1) (f (vfree (Quote ii)))))  

quote ii (VNeutral n)
  =     Inf (neutralQuote ii n)

quote ii VNat       =  Inf Nat
quote ii VZero      =  Zero
quote ii (VSucc n)  =  Succ (quote ii n)

quote ii (VVec a n)         =  Inf (Vec (quote ii a) (quote ii n))
quote ii (VNil a)           =  Nil (quote ii a)
quote ii (VCons a n x xs)   =  Cons  (quote ii a) (quote ii n)
                                        (quote ii x) (quote ii xs)

quote ii (VEq a x y)  =  Inf (Eq (quote ii a) (quote ii x) (quote ii y))
quote ii (VRefl a x)  =  Refl (quote ii a) (quote ii x)

quote ii (VFin n)           =  Inf (Fin (quote ii n))
quote ii (VFZero n)         =  FZero (quote ii n)
quote ii (VFSucc n f)       =  FSucc  (quote ii n) (quote ii f)

neutralQuote :: Int -> Neutral -> ITerm
neutralQuote ii (NFree v)
   =  boundfree ii v
neutralQuote ii (NApp n v)
   =  neutralQuote ii n :$: quote ii v

neutralQuote ii (NNatElim m z s n)
   =  NatElim (quote ii m) (quote ii z) (quote ii s) (Inf (neutralQuote ii n))

neutralQuote ii (NVecElim a m mn mc n xs)
   =  VecElim (quote ii a) (quote ii m)
               (quote ii mn) (quote ii mc)
               (quote ii n) (Inf (neutralQuote ii xs))

neutralQuote ii (NEqElim a m mr x y eq)
   =  EqElim  (quote ii a) (quote ii m) (quote ii mr)
               (quote ii x) (quote ii y)
               (Inf (neutralQuote ii eq))

neutralQuote ii (NFinElim m mz ms n f)
   =  FinElim (quote ii m)
               (quote ii mz) (quote ii ms)
               (quote ii n) (Inf (neutralQuote ii f))

boundfree :: Int -> Name -> ITerm
boundfree ii (Quote k)     =  Bound ((ii - k - 1) `max` 0)
boundfree ii x             =  Free x

instance Show Value where
  show = show . quote0

type Type     =  Value  
type Context  =  [(Name, Type)]

quote0 :: Value -> CTerm
quote0 = quote 0 

lpte :: Ctx Value
lpte =      [(Global "Zero", VNat),
             (Global "Succ", VPi VNat (\ _ -> VNat)),
             (Global "Nat", VStar),
             (Global "natElim", VPi (VPi VNat (\ _ -> VStar)) (\ m ->
                               VPi (m `vapp` VZero) (\ _ ->
                               VPi (VPi VNat (\ k -> VPi (m `vapp` k) (\ _ -> (m `vapp` (VSucc k))))) ( \ _ ->
                               VPi VNat (\ n -> m `vapp` n))))),
             (Global "Nil", VPi VStar (\ a -> VVec a VZero)),
             (Global "Cons", VPi VStar (\ a ->
                            VPi VNat (\ n ->
                            VPi a (\ _ -> VPi (VVec a n) (\ _ -> VVec a (VSucc n)))))),
             (Global "Vec", VPi VStar (\ _ -> VPi VNat (\ _ -> VStar))),
             (Global "vecElim", VPi VStar (\ a ->
                               VPi (VPi VNat (\ n -> VPi (VVec a n) (\ _ -> VStar))) (\ m ->
                               VPi (m `vapp` VZero `vapp` (VNil a)) (\ _ ->
                               VPi (VPi VNat (\ n ->
                                     VPi a (\ x ->
                                     VPi (VVec a n) (\ xs ->
                                     VPi (m `vapp` n `vapp` xs) (\ _ ->
                                     m `vapp` VSucc n `vapp` VCons a n x xs))))) (\ _ ->
                               VPi VNat (\ n ->
                               VPi (VVec a n) (\ xs -> m `vapp` n `vapp` xs))))))),
             (Global "Refl", VPi VStar (\ a -> VPi a (\ x ->
                            VEq a x x))),
             (Global "Eq", VPi VStar (\ a -> VPi a (\ x -> VPi a (\ y -> VStar)))),
             (Global "eqElim", VPi VStar (\ a ->
                              VPi (VPi a (\ x -> VPi a (\ y -> VPi (VEq a x y) (\ _ -> VStar)))) (\ m ->
                              VPi (VPi a (\ x -> m `vapp` x `vapp` x `vapp` VRefl a x)) (\ _ ->
                              VPi a (\ x -> VPi a (\ y ->
                              VPi (VEq a x y) (\ eq ->
                              m `vapp` x `vapp` y `vapp` eq))))))),
             (Global "FZero", VPi VNat (\ n -> VFin (VSucc n))),
             (Global "FSucc", VPi VNat (\ n -> VPi (VFin n) (\ f ->
                             VFin (VSucc n)))),
             (Global "Fin", VPi VNat (\ n -> VStar)),
             (Global "finElim", VPi (VPi VNat (\ n -> VPi (VFin n) (\ _ -> VStar))) (\ m ->
                               VPi (VPi VNat (\ n -> m `vapp` (VSucc n) `vapp` (VFZero n))) (\ _ ->
                               VPi (VPi VNat (\ n -> VPi (VFin n) (\ f -> VPi (m `vapp` n `vapp` f) (\ _ -> m `vapp` (VSucc n) `vapp` (VFSucc n f))))) (\ _ ->
                               VPi VNat (\ n -> VPi (VFin n) (\ f ->
                               m `vapp` n `vapp` f))))))]

lpve :: Ctx Value
lpve =      [(Global "Zero", VZero),
             (Global "Succ", VLam (\ n -> VSucc n)),
             (Global "Nat", VNat),
             (Global "natElim", cEval (Lam (Lam (Lam (Lam (Inf (NatElim (Inf (Bound 3)) (Inf (Bound 2)) (Inf (Bound 1)) (Inf (Bound 0)))))))) ([], [])),
             (Global "Nil", VLam (\ a -> VNil a)),
             (Global "Cons", VLam (\ a -> VLam (\ n -> VLam (\ x -> VLam (\ xs ->
                            VCons a n x xs))))),
             (Global "Vec", VLam (\ a -> VLam (\ n -> VVec a n))),
             (Global "vecElim", cEval (Lam (Lam (Lam (Lam (Lam (Lam (Inf (VecElim (Inf (Bound 5)) (Inf (Bound 4)) (Inf (Bound 3)) (Inf (Bound 2)) (Inf (Bound 1)) (Inf (Bound 0)))))))))) ([],[])),
             (Global "Refl", VLam (\ a -> VLam (\ x -> VRefl a x))),
             (Global "Eq", VLam (\ a -> VLam (\ x -> VLam (\ y -> VEq a x y)))),
             (Global "eqElim", cEval (Lam (Lam (Lam (Lam (Lam (Lam (Inf (EqElim (Inf (Bound 5)) (Inf (Bound 4)) (Inf (Bound 3)) (Inf (Bound 2)) (Inf (Bound 1)) (Inf (Bound 0)))))))))) ([],[])),
             (Global "FZero", VLam (\ n -> VFZero n)),
             (Global "FSucc", VLam (\ n -> VLam (\ f -> VFSucc n f))),
             (Global "Fin", VLam (\ n -> VFin n)),
             (Global "finElim", cEval (Lam (Lam (Lam (Lam (Lam (Inf (FinElim (Inf (Bound 4)) (Inf (Bound 3)) (Inf (Bound 2)) (Inf (Bound 1)) (Inf (Bound 0))))))))) ([],[]))]


iType :: Int -> (NameEnv Value,Context) -> ITerm -> Result (ITerm, Type)
iType ii g (Ann e tyt )
  =     do  tyt' <- cType  ii g tyt VStar
            let ty = cEval tyt (fst g, [])  
            e' <- cType ii g e ty
            return (Ann e tyt, ty)
iType ii g Star   
   =  return (Star, VStar)
-- Any fold instances *within* a type will go away during type compilation,
-- so we need not save those AST tranformations.
iType ii g pi@(Pi tyt tyt')  
   =  do  cType ii g tyt VStar    
          let ty = cEval tyt (fst g, [])  
          cType  (ii + 1) ((\ (d,g) -> (d,  ((Local ii, ty) : g))) g)  
                          (cSubst 0 (Free (Local ii)) tyt') VStar  
          return (pi, VStar)   
iType ii g (Free x)
  =     case lookup x (snd g) of
          Just ty        ->  return (Free x, ty)
          Nothing        ->  throwError ("unknown identifier: " ++ render (iPrint 0 0 (Free x)))
iType ii g (e1 :$: e2)
  =     do  (e1', lhsType) <- iType ii g e1
            case lhsType of
              VPi  ty ty'  ->  do  cType ii g e2 ty
                                   return ( e1' :$: e2, ty' (cEval e2 (fst g, [])))
              _            ->  throwError "illegal application"

iType ii g Nat                  =  return (Nat, VStar)
iType ii g (NatElim m mz ms n)  =
  do  cType ii g m (VPi VNat (const VStar))
      let mVal  = cEval m (fst g, []) 
      mz' <- cType ii g mz (mVal `vapp` VZero)
      ms' <- cType ii g ms (VPi VNat (\ k -> VPi (mVal `vapp` k) (\ _ -> mVal `vapp` VSucc k)))
      n' <- cType ii g n VNat
      let nVal = cEval n (fst g, [])
      return (NatElim (quote ii mVal) mz' ms' n', (mVal `vapp` nVal))

iType ii g (Vec a n) =
  do  a' <- cType ii g a  VStar
      n' <- cType ii g n  VNat
      return (Vec a' n', VStar)
iType ii g (VecElim a m mn mc n vs) =
  do  a' <- cType ii g a VStar
      let aVal = cEval a (fst g, [])
      m' <- cType ii g m
        (  VPi VNat (\n -> VPi (VVec aVal n) (\ _ -> VStar)))
      let mVal = cEval m (fst g, [])
      mn' <- cType ii g mn (foldl vapp mVal [VZero, VNil aVal])
      mc' <- cType ii g mc
        (  VPi VNat (\ n -> 
           VPi aVal (\ y -> 
           VPi (VVec aVal n) (\ ys ->
           VPi (foldl vapp mVal [n, ys]) (\ _ ->
           (foldl vapp mVal [VSucc n, VCons aVal n y ys]))))))
      n' <- cType ii g n VNat
      let nVal = cEval n (fst g, [])
      vs' <- cType ii g vs (VVec aVal nVal)
      let vsVal = cEval vs (fst g, [])
      return (VecElim (quote ii aVal) (quote ii mVal) mn' mc' n' vs',
              foldl vapp mVal [nVal, vsVal])

iType i g (Eq a x y) =
  do  a' <- cType i g a VStar
      let aVal = cEval a (fst g, [])
      x' <- cType i g x aVal
      y' <- cType i g y aVal
      return (Eq a' x' y', VStar)
iType i g (EqElim a m mr x y eq) =
  do  a' <- cType i g a VStar
      let aVal = cEval a (fst g, [])
      m' <- cType i g m
        (VPi aVal (\ x ->
         VPi aVal (\ y ->
         VPi (VEq aVal x y) (\ _ -> VStar)))) 
      let mVal = cEval m (fst g, [])
      mr' <- cType i g mr
        (VPi aVal (\ x ->
         foldl vapp mVal [x, x]))
      x' <- cType i g x aVal
      let xVal = cEval x (fst g, [])
      y' <- cType i g y aVal
      let yVal = cEval y (fst g, [])
      eq' <- cType i g eq (VEq aVal xVal yVal)
      let eqVal = cEval eq (fst g, [])
      return (EqElim (quote i aVal) (quote i mVal) mr' x' y' eq', foldl vapp mVal [xVal, yVal])


cType :: Int -> (NameEnv Value,Context) -> CTerm -> Type -> Result CTerm
cType ii g (Inf e) v 
  =     do  (e', v') <- iType ii g e
            unless ( quote0 v == quote0 v') (throwError ("type mismatch:\n" ++ "type inferred:  " ++ render (cPrint 0 0 (quote0 v')) ++ "\n" ++ "type expected:  " ++ render (cPrint 0 0 (quote0 v)) ++ "\n" ++ "for expression: " ++ render (iPrint 0 0 e)))
            return (Inf e')
cType ii g (Lam e) ( VPi ty ty')
  =     cType  (ii + 1) ((\ (d,g) -> (d,  ((Local ii, ty ) : g))) g)
                (cSubst 0 (Free (Local ii)) e) ( ty' (vfree (Local ii))) 

cType ii g Zero      VNat  =  return Zero
cType ii g (Succ k)  VNat  =  cType ii g k VNat

cType ii g (Nil a) (VVec bVal VZero) =
  do  a' <- cType ii g a VStar
      let aVal = cEval a (fst g, []) 
      unless  (quote0 aVal == quote0 bVal) 
              (throwError "type mismatch")
      return a'
cType ii g (Cons a n x xs) (VVec bVal (VSucc k)) =
  do  a' <- cType ii g a VStar
      let aVal = cEval a (fst g, [])
      unless  (quote0 aVal == quote0 bVal)
              (throwError "type mismatch")
      n' <- cType ii g n VNat
      let nVal = cEval n (fst g, [])
      unless  (quote0 nVal == quote0 k)
              (throwError "number mismatch")
      x' <- cType ii g x aVal
      xs' <- cType ii g xs (VVec bVal k)
      return (Cons a' n' x' xs')

cType ii g (Refl a z) (VEq bVal xVal yVal) =
  do  a' <- cType ii g a VStar
      let aVal = cEval a (fst g, [])
      unless  (quote0 aVal == quote0 bVal)
              (throwError "type mismatch")
      z' <- cType ii g z aVal
      let zVal = cEval z (fst g, [])
      unless  (quote0 zVal == quote0 xVal && quote0 zVal == quote0 yVal)
              (throwError "type mismatch")
      return (Refl a' z')

cType ii g _ _
  =     throwError "type mismatch"
  
{- Printing -} 

vars :: [String]
vars = [ c : n | n <- "" : map show [1..], c <- ['x','y','z'] ++ ['a'..'w'] ]

parensIf :: Bool -> Doc -> Doc
parensIf True  = PP.parens
parensIf False = id

iPrint :: Int -> Int -> ITerm -> Doc
iPrint p ii (Ann c ty)       =  parensIf (p > 1) (cPrint 2 ii c <> text " :: " <> cPrint 0 ii ty)
iPrint p ii Star             =  text "*"
iPrint p ii (Pi d (Inf (Pi d' r)))
                               =  parensIf (p > 0) (nestedForall (ii + 2) [(ii + 1, d'), (ii, d)] r)
iPrint p ii (Pi d r)         =  parensIf (p > 0) (sep [text "forall " <> text (vars !! ii) <> text " :: " <> cPrint 0 ii d <> text " .", cPrint 0 (ii + 1) r])
iPrint p ii (Bound k)        =  text (vars !! (ii - k - 1))
iPrint p ii (Free (Global s))=  text s
iPrint p ii (i :$: c)         =  parensIf (p > 2) (sep [iPrint 2 ii i, nest 2 (cPrint 3 ii c)])
iPrint p ii Nat              =  text "Nat"
iPrint p ii (NatElim m z s n)=  iPrint p ii (Free (Global "natElim") :$: m :$: z :$: s :$: n)
iPrint p ii (Vec a n)        =  iPrint p ii (Free (Global "Vec") :$: a :$: n)
iPrint p ii (VecElim a m mn mc n xs)
                               =  iPrint p ii (Free (Global "vecElim") :$: a :$: m :$: mn :$: mc :$: n :$: xs)
iPrint p ii (Eq a x y)       =  iPrint p ii (Free (Global "Eq") :$: a :$: x :$: y)
iPrint p ii (EqElim a m mr x y eq)
                               =  iPrint p ii (Free (Global "eqElim") :$: a :$: m :$: mr :$: x :$: y :$: eq)
iPrint p ii (Fin n)          =  iPrint p ii (Free (Global "Fin") :$: n)
iPrint p ii (FinElim m mz ms n f)
                               =  iPrint p ii (Free (Global "finElim") :$: m :$: mz :$: ms :$: n :$: f)
iPrint p ii x                 =  text ("[" ++ show x ++ "]")

cPrint :: Int -> Int -> CTerm -> Doc
cPrint p ii (Inf i)    = iPrint p ii i
cPrint p ii (Lam c)    = parensIf (p > 0) (text "\\ " <> text (vars !! ii) <> text " -> " <> cPrint 0 (ii + 1) c)
cPrint p ii Zero       = fromNat 0 ii Zero     --  text "Zero"
cPrint p ii (Succ n)   = fromNat 0 ii (Succ n) --  iPrint p ii (Free (Global "Succ") :$: n)
cPrint p ii (Nil a)    = iPrint p ii (Free (Global "Nil") :$: a)
cPrint p ii (Cons a n x xs) =
                           iPrint p ii (Free (Global "Cons") :$: a :$: n :$: x :$: xs)
cPrint p ii (Refl a x) = iPrint p ii (Free (Global "Refl") :$: a :$: x)
cPrint p ii (FZero n)  = iPrint p ii (Free (Global "FZero") :$: n)
cPrint p ii (FSucc n f)= iPrint p ii (Free (Global "FSucc") :$: n :$: f)

fromNat :: Int -> Int -> CTerm -> Doc
fromNat n ii Zero = int n
fromNat n ii (Succ k) = fromNat (n + 1) ii k
fromNat n ii t = parensIf True (int n <> text " + " <> cPrint 0 ii t)

nestedForall :: Int -> [(Int, CTerm)] -> CTerm -> Doc
nestedForall ii ds (Inf (Pi d r)) = nestedForall (ii + 1) ((ii, d) : ds) r
nestedForall ii ds x                = sep [text "forall " <> sep [parensIf True (text (vars !! n) <> text " :: " <> cPrint 0 n d) | (n,d) <- reverse ds] <> text " .", cPrint 0 ii x]