{-# LANGUAGE GeneralizedNewtypeDeriving, PatternGuards #-}
module Evaluator (
  TC,
  runTC,
  compileTerm,
  constType,
  eval
) where

import qualified Data.Map as Map
import Control.Monad.State hiding (put)
import Control.Monad.Error
import Text.Printf
import Syntax
import Printer
import Parser

-- Type Checker Monad ---------------------------------------------------------

type Env = Map.Map Name Type

newtype TC a = TC (StateT Env (Either String) a)
  deriving (Monad, MonadState Env, MonadError String)

withExtendedEnv :: Name -> Type -> TC a -> TC a
withExtendedEnv k v (TC a) = do
  env <- get
  TC $ lift $ evalStateT a (Map.insert k v env)

runTC :: TC t -> Either String t
runTC (TC a) = evalStateT a (Map.empty)

put :: Name -> Type -> TC ()
put n t = modify (Map.insert n t)

getType :: Name -> TC Type
getType n = do
  env <- get
  case Map.lookup n env of
    Nothing -> throwError $ n ++ " not found in environment"
    Just t  -> return t


-- Constant typing ------------------------------------------------------------

constType :: Constant -> Type
constType c = t
  where
    t = case parseType s of
      Left err -> error $ show err
      Right u  -> u
    s = case c of
      Ctrue     -> "{b: Bool | b}"
      Cfalse    -> "{b: Bool | not b}"
      Cimpl     -> "b1: Bool -> b2: Bool -> {b: Bool | <=> b (<=> b1 b2)}"
      Cimplb b1 -> printf "b2: Bool -> {b: Bool | <=> b (<=> %s b2)}" b1s
        where b1s = if b1 then "true" else "false"
      Cnot      -> "b: Bool -> {b': Bool | <=> b (not b)}"
      Cint n    -> printf "{m: Int | = m %s}" (show n)
      Cplus     -> "n: Int -> m: Int -> {z: Int | = z (+ n m)}"
      Cplusn n  -> printf "m: Int -> {z: Int | = z (+ %s m)}" (show n)
      Ceq       -> "n: Int -> m: Int -> {b: Bool | <=> b (= n m)}"
      Ceqn n    -> printf "m: Int -> {b: Bool | <=> b (= %s m)}" (show n)
      Cdiv      -> "n: Int -> m:{x: Int | not (= x 0)} -> {z: Int | = z (/ n m)}"
      Cdivn n   -> printf "m:{x: Int | not (= x 0)} -> {z: Int | = z (/ %s m)}" (show n)
      Cif u     -> printf "x: Bool -> y:%s -> z:%s -> %s" us us us
        where us = printType True u
      Cfix u    -> printf "x:(y:%s -> %s) -> %s" us us us
        where us = printType True u


-- Compilation of terms -------------------------------------------------------

compileTerm :: Term -> TC (Term,Type)
compileTerm te = case te of

  Var name -> do
    ty <- getType name
    return (te,ty)

  Con con -> return (te, constType con)

  Lam vname vtype body -> do
    vtype' <- compileType vtype
    (body',ty) <- withExtendedEnv vname vtype' (compileTerm body)
    return ((Lam vname vtype' body'), Fun vname vtype' ty)

  App fte bte -> do
    (fte',fty) <- compileTerm fte
    case fty of
      Fun vname vty rty -> do
        bte' <- compileCheckTerm bte vty
        let ty = expandType rty (vname := bte')
        return (App fte' bte', ty)
      _ -> fail $ printf "Non-lambda term application: %s" (pretty te)

  Cast ty1 ty2 tf -> do
    ty1' <- compileType ty1
    ty2' <- compileType ty2
    tf'  <- compileCheckTerm tf ty1'
    return (Cast ty1' ty2' tf', ty2')


-- Checked compilation of terms ------------------------------------------------

compileCheckTerm :: Term -> Type -> TC Term
compileCheckTerm te ty = do
  (te',ty') <- compileTerm te
  case subType ty' ty of
    STProved -> return te'
    STDisproved -> fail $ (pretty ty')++" is not a subtype of "++(pretty ty)
    STUnknown -> return $ Cast ty' ty te'


-- Compilation of types -------------------------------------------------------

compileType :: Type -> TC Type
compileType ty = case ty of

  Fun vname vty rty -> do
    vty' <- compileType vty
    rty' <- withExtendedEnv vname vty' (compileType rty)
    return $ Fun vname vty' rty'

  Ref vname base te -> do
    (te',tz) <- withExtendedEnv vname (mkBaseT base) (compileTerm te)
    case tz of
      Ref _ BBool _ -> return $ Ref vname base te'
      _ -> fail $ "Invalid " ++ pretty base ++ " refinement body: " ++ pretty te

mkBaseT :: Base -> Type
mkBaseT b = Ref "x" b (Con Ctrue)


-- Substition of terms --------------------------------------------------------

data Substitution = Name := Term

expandType :: Type -> Substitution -> Type
expandType ty s@(name := _) = case ty of

  Fun vname vtype rtype
    | vname == name -> ty
    | otherwise -> Fun vname (expandType vtype s) (expandType rtype s)

  Ref vname base te'
    | vname == name -> ty
    | otherwise -> Ref vname base (expandTerm te' s)

expandTerm :: Term -> Substitution -> Term
expandTerm te s@(name := te') = case te of

  Var vname -> if vname == name then te' else te

  Con _ -> te

  Lam vname vtype bte
    | vname == name -> te
    | otherwise -> Lam vname (expandType vtype s) (expandTerm bte s)

  App fte bte -> App (expandTerm fte s) (expandTerm bte s)

  Cast t1 t2 cte -> Cast (expandType t1  s)
                         (expandType t2  s)
                         (expandTerm cte s)


-- Subtyping ------------------------------------------------------------------

data SubTypeResult = STProved | STDisproved | STUnknown

subType :: Type -> Type -> SubTypeResult

subType tsub tsuper | tsub == tsuper = STProved

subType (Fun _ vt1 rt1) (Fun _ vt2 rt2) = case (subType vt2 vt1, subType rt1 rt2) of
  (STProved, STProved) -> STProved
  (STDisproved, _)     -> STDisproved
  (_, STDisproved)     -> STDisproved
  _                    -> STUnknown

subType (Ref _ b1 _) (Ref _ b2 (Con Ctrue)) = if (b1 == b2)
  then STProved
  else STDisproved

subType (Ref v b (App (App (Con Ceq) (Var v')) t@(Con _))) (Ref n b' t')
  | b == b', v == v' = case eval $ expandTerm t' (n := t) of
    Just Ctrue  -> STProved
    Just Cfalse -> STDisproved
    _ -> STUnknown

subType _ _ = STUnknown


-- Evalutation ----------------------------------------------------------------

eval :: Monad m => Term -> m Constant
eval t = do
  case t of
    Con c -> return c
    _ -> step t >>= eval

step :: Monad m => Term -> m Term
step te = case te of

  App (Lam x _ t) s -> return $ expandTerm t (x := s)

  App (Con c) t -> do
    mT <- applyCon c t
    case mT of
      Just t' -> return t'
      Nothing -> step t >>= return . App (Con c)

  Cast (Fun x s1 s2) (Fun _ t1 t2) t
    -> return $ Lam x t1 (Cast s2 t2 (App t (Cast t1 s1 (Var x))))

  Cast m@(Ref _ b1 _) n@(Ref x b2 t) c@(Con _) -> if b1 /= b2
    then castFail
    else do
      c' <- eval (expandTerm t (x := c))
      case c' of
        Ctrue -> return c
        _ -> castFail
    where
      castFail = fail $ printf "Runtime cast failed. %s is not a subtype of %s"
        (pretty m) (pretty n)

  App te1 te2 -> do
    te1' <- step te1
    return $ App te1' te2

  Cast t1@(Ref _ _ _) t2@(Ref _ _ _) tf -> do
    tf' <- step tf
    return $ Cast t1 t2 tf'

  t -> fail $ printf "Subterm %s can't be evaluated further" (pretty t)

applyCon :: Monad m => Constant -> Term -> m (Maybe Term)
applyCon c te = case (c, te) of

  (Cimpl, Con Ctrue)  -> rCon $ Cimplb True
  (Cimpl, Con Cfalse) -> rCon $ Cimplb False

  (Cimplb True, Con Ctrue)   -> rCon Ctrue
  (Cimplb True, Con Cfalse)  -> rCon Cfalse
  (Cimplb False, Con Ctrue)  -> rCon Ctrue
  (Cimplb False, Con Cfalse) -> rCon Ctrue

  (Cnot, Con Ctrue)  -> rCon Cfalse
  (Cnot, Con Cfalse) -> rCon Ctrue

  (Cplus, Con (Cint n)) -> rCon $ Cplusn n

  (Cplusn n, Con (Cint m)) -> rCon $ Cint (n+m)

  (Cdiv, Con (Cint n)) -> rCon $ Cdivn n

  (Cdivn n, Con (Cint m)) -> rCon $ Cint (n `div` m)

  (Ceq, Con (Cint n)) -> rCon $ Ceqn n

  (Ceqn n, Con (Cint m)) -> rCon $ if n == m then Ctrue else Cfalse

  (Cif t, Con Ctrue)  -> return $ Just $ Lam "x" t $ Lam "y" t $ Var "x"
  (Cif t, Con Cfalse) -> return $ Just $ Lam "x" t $ Lam "y" t $ Var "y"

  (Cfix t, tf) -> return $ Just $ App tf $ App (Con $ Cfix t) tf

  (_, Con _) -> fail $ printf "Undefined application: [[%s]] (%s)"
              (pretty c) (pretty te)

  _ -> return Nothing

  where rCon = return . Just . Con