-- | Symbolic Expressons
-- More exercises with recursive data types
-- Functional Programming course 2017.
-- Thomas Hallgren

{-
This started as a skeleton, the definitions were filled in
during the lecture.
-}
--------------------------------------------------------------------------------

module SymbolicExpressions where
import Data.List(union)
import Data.Char(isLetter,isDigit)
import Parsing

-- | A Haskell data type for arithmetic expressions with variables
data Expr = Num Integer
          | Var Name
          | Add Expr Expr
          | Mul Expr Expr
          deriving (Eq)

type Name = String

ex1 = Num 2
ex2 = Add (Num 2) (Num 2)
ex3 = Mul (Add (Num 1) (Num 2)) (Num 3)
ex4 = Add (Num 1) (Mul (Num 2) (Num 3))
ex5 = Add (Mul (Num 2) x) (Mul (Num 3) y)
ex6 = Add (Mul (Num 2) (Mul x x)) (Mul (Num 3) y)

x   = Var "x"
y   = Var "y"

--------------------------------------------------------------------------------
-- * Showing expressions

instance Show Expr where
  show = showExpr

showExpr :: Expr -> String
showExpr (Add e1 e2) = showExpr e1 ++ "+" ++ showExpr e2
showExpr e           = showFactor e

showFactor :: Expr -> String
showFactor (Mul e1 e2) = showFactor e1 ++ "*" ++ showFactor e2
showFactor (Num n) = show n
showFactor (Var x) = x
showFactor e       = "("++showExpr e++")"

--------------------------------------------------------------------------------
-- * Parsing expressions

expr, term, factor :: Parser Expr

expr = leftAssoc Add term (char '+')
term = leftAssoc Mul factor (char '*')

factor =   Num <$> number
       <|> Var <$> name
       <|> char '(' *> expr <* char ')'

name :: Parser Name
name = oneOrMore (sat isLetter)

number :: Parser Integer
number = read <$> oneOrMore (sat isDigit)

leftAssoc :: (t->t->t) -> Parser t -> Parser sep -> Parser t
leftAssoc op item sep = foldl1 op <$> chain item sep

--------------------------------------------------------------------------------

-- | Gathering variables

vars :: Expr -> [Name]
vars (Num n) = []
vars (Var x) = [x]
vars (Add a b) = vars a `union` vars b
vars (Mul a b) = vars a `union` vars b



-- | Substituting expressions for variables
substitute :: [(Name,Expr)] -> Expr -> Expr
substitute env (Num n) = Num n
substitute env (Var x) = case lookup x env of
                           Nothing -> Var x
                           Just e  -> e
substitute env (Add a b) = Add (substitute env a) (substitute env b)
substitute env (Mul a b) = Mul (substitute env a) (substitute env b)


--------------------------------------------------------------------------------

-- | Evaluating Symbolic Expressions

eval :: [(Name,Integer)] -> Expr -> Integer
eval env (Num n)   = n
eval env (Var x)   = case lookup x env of
                       Just n  -> n
                       Nothing -> error ("undefined variable: "++x)
eval env (Add a b) = eval env a + eval env b
eval env (Mul a b) = eval env a * eval env b


--------------------------------------------------------------------------------
-- * Symbolic Differentiation

-- | Symbolic Differentiation function
diff :: Expr -> Name -> Expr
diff (Num n)   x = Num 0
diff (Var y)   x | y==x      = Num 1
                 | otherwise = Num 0
diff (Add a b) x = add (diff a x) (diff b x)
diff (Mul a b) x = add (mul a (diff b x)) (mul (diff a x) b)

add (Num 0) b = b
add a (Num 0) = a
add (Num a) (Num b) = Num (a+b)
add a b | a==b = Mul (Num 2) a
add a b = Add a b

mul (Num 0) b = Num 0
mul a (Num 0) = Num 0
mul (Num 1) b = b
mul a (Num 1) = a
mul (Num a) (Num b) = Num (a*b)
mul (Num a) (Mul (Num b) c) = Mul (Num (a*b)) c
mul a b = Mul a b

--------------------------------------------------------------------------------
-- * A simple calculator with definitions

data Command = Define Name Expr | Eval Expr

command :: Parser Command
command =   Define <$> name <* char '=' <*> expr
        <|> Eval   <$> expr

main = do putStrLn "Welcome to the simple calculator!"
          readEvalPrintLoop []

-- | Our read-eval-print loop
readEvalPrintLoop env =
  do putStr "Expression? "
     s <- getLine
     case completeParse command s of
       Just (Define x e) -> readEvalPrintLoop ((x,eval env e):env)
       Just (Eval e)     -> do print (eval env e)
                               readEvalPrintLoop env
       _ -> putStrLn "Syntax error!"

--------------------------------------------------------------------------------

completeParse p s =
  case parse p s of
    Just (e,"") -> Just e
    _ -> Nothing

{- Below are the tests we ran in GHCi:
:l SymbolicExpressions.hs
:r
ex5
ex6
:r
vars ex5
vars ex6
ex3
vars ex3
:t lookup
:r
substitute [] ex5
substitute [] ex6
substitute [("x",Num 5)] ex6
substitute [("x",Num 5),("y",ex3)] ex6
:r
ex3
eval [] ex3
ex5
eval [(] ex5
eval [] ex5
eval [("y",2)] ex5
eval [("y",2),("x",5)] ex5
:r
ex2
diff ex2 "x"
ex3
diff ex3 "x"
ex5
diff ex5 "x"
:r
diff ex5 "x"
:r
diff ex5 "x"
:r
diff ex5 "x"
ex4
ex5
ex6
diff ex6 "x"
:r
ex6
diff ex6 "x"
:r
diff ex6 "x"
:r
diff ex6 "x"
ex6
:r
main
:r
main
:r
main
main
:r
main
-}