module LP where
import Control.Monad.Error
import Data.List
import Data.Char
import Text.PrettyPrint.HughesPJ hiding (parens)
import qualified Text.PrettyPrint.HughesPJ as PP

--------------------------------------------------------------------------------
------------------------------------ Types -------------------------------------
--------------------------------------------------------------------------------

type NameEnv v = [(Name, v)]
 
-- Besides numbers, also allow strings to name global constants
data Name
   =  Global  String
   |  Local   Int
   |  Quote   Int
  deriving (Show, Eq)

type Result a = Either String a

-- Checkable terms
data CTerm
   =  Inf ITerm
   |  Lam CTerm
   |  Zero
   |  Succ CTerm
   |  Nil CTerm
   |  Cons CTerm CTerm CTerm CTerm

  deriving (Show, Eq)

-- Inferable terms
data ITerm
   =  Ann CTerm CTerm
   |  Star
   |  Pi CTerm CTerm
   |  Bound  Int
   |  Free Name
   |  ITerm :$: CTerm
   |  Nat
   |  NatElim CTerm CTerm CTerm CTerm
   |  Vec CTerm CTerm
   |  VecElim CTerm CTerm CTerm CTerm CTerm CTerm
  deriving (Show, Eq)

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

-- Stuck values, aka neutral terms
data Neutral
   =  NFree Name
   |  NApp Neutral Value
   |  NNatElim Value Value Value Neutral
   |  NVecElim Value Value Value Value Value Neutral

--------------------------------------------------------------------------------
----------------------------------  Evaluation ---------------------------------
--------------------------------------------------------------------------------

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)

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)

--------------------------------------------------------------------------------
-------------------------------- Substitution ----------------------------------
--------------------------------------------------------------------------------


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)

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)

--------------------------------------------------------------------------------
--------------------------------- Quotation ------------------------------------
--------------------------------------------------------------------------------

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)

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

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 

--------------------------------------------------------------------------------
------------------------------- Type checking ----------------------------------
--------------------------------------------------------------------------------

iType0 :: (NameEnv Value,Context) -> ITerm -> Result Type
iType0 = iType 0

iType :: Int -> (NameEnv Value,Context) -> ITerm -> Result Type
iType ii g (Ann e tyt )
  =     do  cType  ii g tyt VStar
            let ty = cEval tyt (fst g, [])  
            cType ii g e ty
            return ty
iType ii g Star   
   =  return VStar   
iType ii g (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 VStar   
iType ii g (Free x)
  =     case lookup x (snd g) of
          Just ty        ->  return ty
          Nothing        ->  throwError ("unknown identifier: " ++ render (iPrint 0 0 (Free x)))
iType ii g (e1 :$: e2)
  =     do  si <- iType ii g e1
            case si of
              VPi  ty ty'  ->  do cType ii g e2 ty
                                  return ( ty' (cEval e2 (fst g, [])))
              _            ->  throwError "illegal application"
iType ii g Nat                  =  return VStar
iType ii g (NatElim m mz ms n)  =
  do  cType ii g m (VPi VNat (const VStar))
      let mVal  = cEval m (fst g, []) 
      cType ii g mz (mVal `vapp` VZero)
      cType ii g ms (VPi VNat (\ k -> VPi (mVal `vapp` k) (\_ -> mVal `vapp` VSucc k)))
      cType ii g n VNat
      let nVal = cEval n (fst g, [])
      return (mVal `vapp` nVal)
iType ii g (Vec a n) =
  do  cType ii g a  VStar
      cType ii g n  VNat
      return VStar
iType ii g (VecElim a m mn mc n vs) =
  do  cType ii g a VStar
      let aVal = cEval a (fst g, [])
      cType ii g m
        (  VPi VNat (\n -> VPi (VVec aVal n) (\_ -> VStar)))
      let mVal = cEval m (fst g, [])
      cType ii g mn (foldl vapp mVal [VZero, VNil aVal])
      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]))))))
      cType ii g n VNat
      let nVal = cEval n (fst g, [])
      cType ii g vs (VVec aVal nVal)
      let vsVal = cEval vs (fst g, [])
      return (foldl vapp mVal [nVal, vsVal])

cType :: Int -> (NameEnv Value,Context) -> CTerm -> Type -> Result ()
cType ii g (Inf e) v 
  =     do  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)))
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 ()
cType ii g (Succ k)  VNat  =  cType ii g k VNat
cType ii g (Nil a) (VVec bVal VZero) =
  do  cType ii g a VStar
      let aVal = cEval a (fst g, []) 
      unless  (quote0 aVal == quote0 bVal) 
              (throwError "type mismatch")
cType ii g (Cons a n x xs) (VVec bVal (VSucc k)) =
  do  cType ii g a VStar
      let aVal = cEval a (fst g, [])
      unless  (quote0 aVal == quote0 bVal)
              (throwError "type mismatch")
      cType ii g n VNat
      let nVal = cEval n (fst g, [])
      unless  (quote0 nVal == quote0 k)
              (throwError "number mismatch")
      cType ii g x aVal
      cType ii g xs (VVec bVal k)
cType ii g  _ _
  =     throwError "type mismatch"

--------------------------------------------------------------------------------
-------------------------------- Pretty printing -------------------------------
--------------------------------------------------------------------------------

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

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

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]