module Main where
import Text.Parsec hiding (State, token)
import Text.Parsec.String
import Text.Parsec.Expr
import Control.Applicative((<$>), (<*>))
import Control.Monad.State
import System.Console.Haskeline
import System.Environment (getArgs)
import Data.Fixed (mod')
import Data.List ((\\))
import Data.Set (Set, empty, singleton, delete, union, member, toList)

-- The abstract syntax
type Id = String

-- Binary operators
data Op = Add | Sub | Mul | Div | Mod
        | Lt | Gt | Le | Ge
        | Eq | SEq | Ne | SNe
        | And | Or
  deriving (Show, Eq)

-- Unary operators
data UOp = Minus | Not deriving (Show, Eq)
           
-- Expressions
data Exp = Un UOp Exp         {- A unary operation -}
         | Bin Op Exp Exp     {- A binary operation -}
         | Tern Exp Exp Exp   {- e0 ? e1 : e2 -}
         | Undefined          {- the undefined constant -}
         | K Double           {- numbers -}
         | B Bool             {- booleans -}
         | Fun Id Exp
         | App Exp Exp
         | Var Id
         | Assign Id Exp
         | Seq Exp Exp  
         | Closure Loc Id Exp Store
  deriving (Eq)

-- Loc used to label closures with a "location".
-- This simulates the memory address of the closure in JavaScript
-- and is needed to implement ==.
type Loc = Int

-- Stores (mutable environments)
type Store = [(Id, Exp)]
type Stack = [Store]

lookupVar :: Id -> Store -> Exp
lookupVar x sigma = maybe Undefined id $ lookup x sigma

setVar :: Store -> Id -> Exp -> Store
setVar sigma x v = (x,v):sigma

-- Memory location (used to distinguish between different closures)
newLoc :: State (Loc, Stack) Loc
newLoc = do
           (nextLoc, stack) <- get
           put (nextLoc + 1, stack)
           return nextLoc

-- Stack
emptyStack :: Stack
emptyStack = [[]]

getStore :: State (Loc, Stack) Store
getStore = do
           (nextLoc, sigma:stackTail) <- get
           return sigma

setStore :: Id -> Exp -> State (Loc, Stack) ()
setStore x v = do
           (nextLoc, sigma:stackTail) <- get
           put (nextLoc, (setVar sigma x v):stackTail)
           return ()

pushStack :: Store -> State (Loc, Stack) ()
pushStack sigma = do
             (nextLoc, stack) <- get
             put (nextLoc, sigma:stack)
             return ()

popStack :: State (Loc, Stack) ()
popStack = do
             (nextLoc, sigma:stackTail) <- get
             put (nextLoc, stackTail)
             return ()

-- The evaluator
step :: Exp -> State (Loc, Stack) Exp

-- Values
step e | value e = do
                    popStack
                    return e

-- Make closure
step (Fun x e)  = do
                    loc <- newLoc
                    sigma <- getStore
                    return $ Closure loc x e sigma

-- Seq congruence
step (Seq x y)
  | not (value x)            = do
                                x' <- step x
                                return $ Seq x' y

-- Seq computation
step (Seq x y)               = return y

-- Assign congruence
step (Assign x e)
  | not (value e)            = do
                                e' <- step e
                                return $ Assign x e'

-- Assign computation
step (Assign x v)            = do
                                setStore x v
                                return v

-- Apply congruence
step (App x y)
  | not (value x)            = do
                                x' <- step x
                                return $ App x' y
  | value x && not (value y) = do
                                y' <- step y
                                return $ App x y'

-- Apply computation
step (App (Closure loc x e sigma) v) = do
                     pushStack (setVar sigma x v)
                     return e

-- Var
step (Var x) = do
                sigma <- getStore
                return $ lookupVar x sigma

-- Unary congruence
step (Un op x)
  | not (value x) = do 
                      x' <- step x
                      return $ Un op x'

-- Unary computation
step (Un Minus e) = return $ K (-(toNumber e))
step (Un Not e)   = return $ B (not (toBool e))

-- Ternary congruence
step (Tern c e1 e2)
  | not (value c) = do
                      c' <- step c
                      return $ Tern c' e1 e2

-- Ternary computation
step (Tern c e1 e2)
  | toBool c  = return e1
  | otherwise = return e2

-- First handle && and || where the first operand is a value.
-- We need to do this before the usual congruence rules to get short-circuiting right.
step (Bin And x y)
  | value x = return $ if toBool x then y else x

step (Bin Or x y)
  | value x = return $ if toBool x then x else y

-- Binary congruence
step (Bin op x y)
  | not (value x)            = do
                                x' <- step x
                                return $ Bin op x' y
  | value x && not (value y) = do
                                y' <- step y
                                return $ Bin op x y'

-- Binary computation...

-- arithmetic
step (Bin Add x y) = return $ K (toNumber x + toNumber y)
step (Bin Sub x y) = return $ K (toNumber x - toNumber y)
step (Bin Mul x y) = return $ K (toNumber x * toNumber y)
step (Bin Div x y) = return $ K (toNumber x / toNumber y)
step (Bin Mod x y) = return $ K (toNumber x `rem'` toNumber y)
  where
    rem' :: Double -> Double -> Double
    rem' x y
      | isNaN x        = nan
      | isNaN y        = nan
      | isInfinite x   = nan
      | y == 0         = nan
      | isInfinite y   = x
      | x < 0 && y > 0 = 0 - mod' (-x) y
      | x > 0 && y < 0 = mod' x (-y)
      | otherwise      = mod' x y

-- relational
step (Bin Lt x y)  = return $ B (toNumber x < toNumber y)
step (Bin Gt x y)  = return $ B (toNumber x > toNumber y)
step (Bin Le x y)  = return $ B (toNumber x <= toNumber y)
step (Bin Ge x y)  = return $ B (toNumber x >= toNumber y)

-- strict equality and disequality
step (Bin SEq (B x) (B y)) = return $ B (x == y)
step (Bin SEq (K x) (K y)) = return $ B (x == y)
step (Bin SEq (K x) (B y)) = return $ B False
step (Bin SEq (B x) (K y)) = return $ B False
step (Bin SEq Undefined Undefined) = return $ B True
step (Bin SEq (Closure x _ _ _) (Closure y _ _ _)) = return $ B (x == y)
step (Bin SEq v1 v2) = return $ B False

step (Bin SNe (B x) (B y)) = return $ B (x /= y)
step (Bin SNe (K x) (K y)) = return $ B (x /= y)
step (Bin SNe (K x) (B y)) = return $ B True
step (Bin SNe (B x) (K y)) = return $ B True
step (Bin SNe Undefined Undefined) = return $ B False
step (Bin SNe (Closure x _ _ _) (Closure y _ _ _)) = return $ B (x /= y)
step (Bin SNe v1 v2) = return $ B True

-- non-strict equality and disequality
step (Bin Eq (B x) (B y)) = return $ B (x == y)
step (Bin Eq (K x) (K y)) = return $ B (x == y)
step (Bin Eq (K x) (B y)) = return $ B (toNumber (K x) == toNumber (B y))
step (Bin Eq (B x) (K y)) = return $ B (toNumber (B x) == toNumber (K y))
step (Bin Eq Undefined Undefined) = return $ B True
step (Bin Eq (Closure x _ _ _) (Closure y _ _ _)) = return $ B (x == y)
step (Bin Eq v1 v2) = return $ B False

step (Bin Ne (B x) (B y)) = return $ B (x /= y)
step (Bin Ne (K x) (K y)) = return $ B (x /= y)
step (Bin Ne (K x) (B y)) = return $ B (toNumber (K x) /= toNumber (B y))
step (Bin Ne (B x) (K y)) = return $ B (toNumber (B x) /= toNumber (K y))
step (Bin Ne Undefined Undefined) = return $ B False
step (Bin Ne (Closure x _ _ _) (Closure y _ _ _)) = return $ B (x /= y)
step (Bin Ne v1 v2) = return $ B True

-- values
step v @ (K _)             = return v
step v @ (B _)             = return v
step v @ Undefined         = return v
step v @ (Closure _ _ _ _) = return v

-- catch missing patterns
step e = error $ "missing pattern " ++ (show e)

-- is this term a value?
value (K _) = True
value (B _) = True
value Undefined = True
value (Closure _ _ _ _) = True
value _     = False

-- evaluate an expression given a step function, returning the trace of
-- intermediate terms
-- stop stepping when the term does not change
trace :: Stack -> Exp -> ([Exp], (Loc, Stack))
trace stack e = runState (traceM e) (0, stack)
  where
    traceM :: Exp -> State (Loc, Stack) [Exp]
    traceM e = if value e
                  then return [e]
                  else do
                          e' <- step e
                          es <- traceM e'
                          return (e:es)


-- evaluate an expression given a step function
eval :: Stack -> Exp -> (Exp, (Loc, Stack))
eval stack e = runState (evalM e) (0, stack)
  where
    evalM :: Exp -> State (Loc, Stack) Exp
    evalM e = if not (value e) then
                do 
                  e' <- step e
                  evalM e'
              else return e

-- name checker
-- check the top-level term is closed
checkNames :: Monad m => Exp -> (Id -> m ()) -> m () -> m ()
checkNames e printLn performEval =
  case toList $ fv e of
    [] -> performEval
    xs -> mapM_ printLn $ map (\x -> "variable " ++ x ++ " not found") xs

-- Free variables
-- Return the free variables in the given term.
fv :: Exp -> Set Id

fv (Var x)        = singleton x
fv (App e1 e2)    = fv e1 `union` fv e2
fv (Fun x e)      = delete x (fv e)

fv (Seq e1 e2)    = fv e1 `union` fv e2
fv (Assign x e)   = singleton x `union` fv e

fv (Un op e)      = fv e
fv (Bin op e1 e2) = fv e1 `union` fv e2
fv (Tern c e1 e2) = fv c `union` fv e1 `union` fv e2
fv (K n)          = empty
fv (B b)          = empty
fv Undefined      = empty
fv e              = error $ "missing case: no fv for " ++ (show e)

-- main and the REPL
main :: IO ()
main = do
  args <- getArgs
  if null args then
    runInputT defaultSettings readEvalPrintLoop
  else do
    input <- readFile $ head args
    ok <- replCmd putStr root input
    return ()

replCmd :: Monad m => (String -> m ()) -> Parser Cmd -> String -> m Bool
replCmd outputStr parser line = case parse parser "input" line of
  Right (Trace e) -> do
    checkNames e outputStrLn (outputStr $ unlines $ (map show . fst) $ trace emptyStack e)
    return True
  Right (Eval e)  -> do
    checkNames e outputStrLn (outputStrLn $ (show . fst) $ eval emptyStack e)
    return True
  Right Quit ->
    return False
  Left err -> do
    outputStrLn $ show err
    return True
  where
    outputStrLn = outputStr . (++ "\n")

readEvalPrintLoop :: InputT IO ()
readEvalPrintLoop = do
  maybeLine <- getInputLine "> "
  case maybeLine of
    Nothing -> return ()
    Just line -> do
      ok <- replCmd outputStr command line
      if ok then readEvalPrintLoop else return ()

-- REPL commands
data Cmd = Trace Exp | Eval Exp | Quit

-- A simple parser
command :: Parser Cmd
command = (do
    spaces
    c <- (do { char ':'; spaces; (trace <|> quit)}) <|> eval
    spaces
    eof
    return c) <?> "expression, :trace expression, or :quit"
  where
    eval  = do { x <- expr; return (Eval x) }
    trace = do { string "trace"; spaces; x <- expr; return (Trace x) }
    quit  = do { string "quit"; return Quit }

root :: Parser Cmd
root = (do
    spaces
    e <- expr
    eof
    return $ Eval e) <?> "top-level expression"

-- expression parser
expr :: Parser Exp
expr = (try fun) <|> cond <?> "expression"

-- Ensure all 'char' and 'string' are followed by 'spaces'.
-- These should be the only places 'spaces' are needed.
token :: String -> Parser ()
token s = do { string s; spaces; return () } <?> s

tchar :: Char -> Parser ()
tchar c = do { char c; spaces; return () } <?> (c:"")

fun :: Parser Exp
fun = (do
  token "function"
  tchar '('
  x <- ident
  tchar ')'
  tchar '{'
  e <- body
  option () (do { tchar ';'; return ()})
  tchar '}'
  return $ Fun x e
  ) <?> "function"

ret :: Parser Exp
ret = do { token "return"; expr } <?> "return expression"

body :: Parser Exp
body = do
  ss <- many (do { a <- try assign; tchar ';'; return a })
  e <- ret
  return $ foldr Seq e ss

assign :: Parser Exp
assign = do
  x <- ident
  tchar '='
  e <- expr
  return $ Assign x e

cond :: Parser Exp
cond = do
  c <- try binary
  option c (do
    tchar '?'
    t <- expr
    tchar ':'
    e <- expr
    return $ Tern c t e)

binary :: Parser Exp
binary = buildExpressionParser table factor <?> "expression" where
  table = 
    [
    [Prefix (do { tchar '-'; return $ Un Minus })],
    [Prefix (do { tchar '!'; return $ Un Not })],
    [Infix (do { tchar '*'; return $ Bin Mul }) AssocLeft,
     Infix (do { tchar '/'; return $ Bin Div }) AssocLeft,
     Infix (do { tchar '%'; return $ Bin Mod }) AssocLeft],
    [Infix (do { tchar '+'; return $ Bin Add }) AssocLeft,
     Infix (do { tchar '-'; return $ Bin Sub }) AssocLeft],
    [Infix (do { try $ token ">="; return $ Bin Ge }) AssocLeft,
     Infix (do { tchar '>'; return $ Bin Gt }) AssocLeft,
     Infix (do { try $ token "<="; return $ Bin Le }) AssocLeft,
     Infix (do { tchar '<'; return $ Bin Lt }) AssocLeft],
    [Infix (do { try $ token "==="; return $ Bin SEq }) AssocLeft,
     Infix (do { token "=="; return $ Bin Eq }) AssocLeft,
     Infix (do { try $ token "!=="; return $ Bin SNe }) AssocLeft,
     Infix (do { try $ token "!="; return $ Bin Ne }) AssocLeft],
    [Infix (do { token "&&"; return $ Bin And }) AssocLeft],
    [Infix (do { token "||"; return $ Bin Or }) AssocLeft]
    ]

par :: Parser Exp
par = (do { tchar '('; x <- expr; tchar ')'; return x })
        <?> "parenthesized expression"

call :: Parser Exp
call = do
  f <- var <|> par
  es <- many par
  return $ foldl App f es 

factor :: Parser Exp
factor = (do 
            f <- number <|> try boolean <|> try undef <|> call
            return f) <?> "simple expression"

sign :: Parser Char
sign = char '+' <|> char '-'

makeNumber '+' ds [] = ds
makeNumber '+' ds mds = ds ++ "." ++ mds
makeNumber '-' ds [] = "-" ++ ds
makeNumber '-' ds mds = "-" ++ ds ++ "." ++ mds

number :: Parser Exp
number = (K . read) <$> (do { s <- option '+' sign;
                              ds <- many1 digit;
                              mds <- option "0" $ do { char '.'; ds' <- many digit; return ds'};
                              spaces;
                              return $ makeNumber s ds mds})

keywords = ["true", "false", "undefined", "function", "return"]

ident :: Parser Id
ident = do 
  c <- letter
  s <- many (letter <|> digit)
  spaces
  let x = c:s in if x `elem` keywords then fail "couldn't match identifier" else return x

var :: Parser Exp
var = Var <$> ident

undef :: Parser Exp
undef = do {token "undefined"; return Undefined}

boolean :: Parser Exp
boolean =  B <$> (do {token "true"; return True} <|> do {token "false"; return False})

-- A simple pretty-printer for Exps
instance Show Exp where
  show (K n) = show n
  show (B True) = "true"
  show (B False) = "false"
  show Undefined = "undefined"
  show e @ (Bin Add e1 e2) = (paren e1) ++ " + " ++ (paren e2)
  show e @ (Bin Sub e1 e2) = (paren e1) ++ " - " ++ (paren e2)
  show e @ (Bin Mul e1 e2) = (paren e1) ++ " * " ++ (paren e2)
  show e @ (Bin Div e1 e2) = (paren e1) ++ " / " ++ (paren e2)
  show e @ (Bin Mod e1 e2) = (paren e1) ++ " % " ++ (paren e2)
  show e @ (Bin And e1 e2) = (paren e1) ++ " && " ++ (paren e2)
  show e @ (Bin Or e1 e2) = (paren e1) ++ " || " ++ (paren e2)
  show e @ (Bin Lt e1 e2) = (paren e1) ++ " < " ++ (paren e2)
  show e @ (Bin Gt e1 e2) = (paren e1) ++ " > " ++ (paren e2)
  show e @ (Bin Le e1 e2) = (paren e1) ++ " <= " ++ (paren e2)
  show e @ (Bin Ge e1 e2) = (paren e1) ++ " >= " ++ (paren e2)
  show e @ (Bin Eq e1 e2) = (paren e1) ++ " == " ++ (paren e2)
  show e @ (Bin SEq e1 e2) = (paren e1) ++ " === " ++ (paren e2)
  show e @ (Bin Ne e1 e2) = (paren e1) ++ " != " ++ (paren e2)
  show e @ (Bin SNe e1 e2) = (paren e1) ++ " !== " ++ (paren e2)
  show (Tern c t e) = (paren c) ++ " ? " ++ (paren t) ++ " : " ++ (paren e)
  show (Un Minus e) = "-" ++ (paren e)
  show (Un Not e) = "!" ++ (paren e)
  show (Fun x e) = "function (" ++ x ++ ") { return " ++ (show e) ++ "; }"
  show (Closure _ x e _) = "function (" ++ x ++ ") { return " ++ (show e) ++ "; }"
  show (App e1 e2) = (paren e1) ++ "(" ++ (paren e2) ++ ")"
  show (Var x) = x
  show (Seq e1 e2) = (paren e1) ++ "; " ++ (paren e2)
  show (Assign x e) = x ++ " = " ++ (show e)

paren :: Exp -> String
paren e @ (Bin _ _ _) = "(" ++ show e ++ ")"
paren e @ (Fun _ _) = "(" ++ show e ++ ")"
paren e @ (App _ _) = "(" ++ show e ++ ")"
paren e @ _           = show e

-- Utility functions used during evaluation
toNumber :: Exp -> Double
toNumber (B False) = 0
toNumber (B True) = 1
toNumber (K n) = n
toNumber _ = nan

toBool :: Exp -> Bool
toBool (K n) | n == 0       = False
             | isNaN n      = False
             | isInfinite n = True
             | otherwise    = True
toBool (B b) = b
toBool (Fun _ _) = True
toBool e = False

nan :: Double
nan = 0/0

inf :: Double
inf = 1/0

minf :: Double
minf = (-1)/0