import Control.Applicative
import Control.Monad

import qualified Text.Parsec as P

data Prog = Let String Expr
          | Expr Expr
          deriving (Show)

data Expr = Num Double
          | Var String
          | App Expr Expr
          | Lam String Expr
          | BinOp Oper Expr Expr
          deriving (Show)

data Value  = VNum Double
            | VFun (Value -> Value)

instance Show Value where
  show (VNum n) = show n
  show (VFun f) = "<func>"

numval :: Value -> Double
numval (VNum n) = n
numval v        = error ("not a number: " ++ show v)

liftFun1 :: (Double -> Double) -> Value
liftFun1 f = VFun $ \x -> VNum (f (numval x))

liftFun2 :: (Double -> Double -> Double) -> Value
liftFun2 f = VFun $ \x -> liftFun1 (\m -> f (numval x) m)

liftFun3 :: (Double -> Double -> Double -> Double) -> Value
liftFun3 f = VFun $ \x -> liftFun2 (\m o -> f (numval x) m o)

data Oper = Add | Sub | Mul | Div
          deriving (Show)

type TokenParser a = P.Parsec String () a

lparen, rparen :: TokenParser ()
lparen = void (P.char '(')
rparen = void (P.char ')')

parens :: TokenParser Expr -> TokenParser Expr
parens p = lparen *> p <* rparen

expr :: TokenParser Expr
expr = lam <|> expr1
  where
    expr1 = do
      v <- expr2
      P.option v (opr fopr1 expr1 v)
    
    expr2 = do
      v <- expr3
      P.option v (opr fopr2 expr2 v)
    
    -- application
    expr3 = do
      v <- expr4
      go v
      where
        go v = do
          P.option v (go' v)
        go' v = do
          P.char ' '
          w <- expr4
          go (App v w)
    
    expr4 = parens expr <|> expr5
    
    expr5 = num <|> var

num :: TokenParser Expr
num = Num <$> fmap read (P.many1 P.digit)

lam :: TokenParser Expr
lam = do
  P.char '\\'
  v <- P.many1 (P.oneOf ['a'..'z'])
  P.string "->"
  e <- expr
  return (Lam v e)

var :: TokenParser Expr
var = Var <$> P.many1 (P.oneOf ['a'..'z'])

opr :: TokenParser (Expr -> Expr -> Expr) -> TokenParser Expr -> Expr -> TokenParser Expr
opr popr prhs lhs = do  { f <- popr
                        ; rhs <- prhs
                        ; return (f lhs rhs)
                        }

fopr1 :: TokenParser (Expr -> Expr -> Expr)
fopr1 = P.choice  [ P.char '+' *> pure (BinOp Add)
                  , P.char '-' *> pure (BinOp Sub)
                  ]

fopr2 :: TokenParser (Expr -> Expr -> Expr)
fopr2 = P.choice  [ P.char '*' *> pure (BinOp Mul)
                  , P.char '/' *> pure (BinOp Div)
                  ]

prog :: TokenParser [Prog]
prog = P.many ((plet <|> pexpr) <* P.newline)
  where
    plet = do
      P.string "let "
      v <- P.many1 (P.oneOf ['a'..'z'])
      P.string " = "
      e <- expr
      return (Let v e)
    
    pexpr = Expr <$> expr

eval :: [(String, Value)] -> Expr -> Value
eval env (Num n) = VNum n
eval env (Lam v e) = VFun $ \x -> eval ((v, x):env) e
eval env (Var k) =
  case lookup k env of
    Just v    -> v
    otherwise -> error ("undefined variable: " ++ k)
eval env (App f e) =
  case eval env f of
    VFun f'   -> f' (eval env e)
    v         -> error ("value not callable: " ++ show v)
eval env (BinOp Add lhs rhs) = VNum $ eval' env lhs + eval' env rhs
eval env (BinOp Sub lhs rhs) = VNum $ eval' env lhs - eval' env rhs
eval env (BinOp Mul lhs rhs) = VNum $ eval' env lhs * eval' env rhs
eval env (BinOp Div lhs rhs) = VNum $ eval' env lhs / eval' env rhs

eval' :: [(String, Value)] -> Expr -> Double
eval' env e = numval (eval env e)

run :: [(String, Value)] -> [Prog] -> [Value]
run env [] = []
run env ((Let v e):ps) = run ((v, eval env e):env) ps
run env ((Expr e):ps) = eval env e : run env ps

main = do
  txt <- readFile "expr.in"
  let mex = P.runParser prog () "<stdin>" txt
  let env = [("sqrt", liftFun1 sqrt), ("ifz", liftFun3 ifz)]
  case mex of
    Right ex  -> print ex *> print (run env ex)
    Left err -> print err
  where
    ifz b x y = if b == 0 then x else y