-- | Symbolic Expressons
-- Exercises with recursive data types and monads
-- Functional Programming course 2016.
-- Thomas Hallgren

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

module SymbolicExpressions where
import Data.List(union)
import Data.Maybe(fromJust)
import Control.Monad(ap,liftM)

-- | A Haskell data type for arithmetic expressions with variables
data Expr = Num Integer
          | Var Name           -- new
          | Add Expr Expr
          | Mul Expr Expr
          | Div Expr Expr      -- new
          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 = Div x (Mul (Num 2) y)
ex6 = Mul (Num 2) x

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

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

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 (Div e1 e2) = showFactor e1 ++ "/" ++ showDivisor e2
showFactor e           = showDivisor e

showDivisor :: Expr -> String
showDivisor (Num n) = show n
showDivisor (Var x) = x
showDivisor e       = "("++showExpr e++")"


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

-- | 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
vars (Div a b) = vars a `union` vars b


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

-- | Evaluating Symbolic Expressions

eval :: [(Name,Integer)] -> Expr -> Integer
eval env (Num n) = n
eval env (Var x) = fromJust (lookup x env)
eval env (Add a b) = eval env a + eval env b
eval env (Mul a b) = eval env a * eval env b
eval env (Div a b) = eval env a `div` eval env b


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

-- | Symbolic Differentiation

diff :: Expr -> Name -> Expr
diff (Num n)   x = Num 0
diff (Var y)   x = if y==x then Num 1 else Num 0
diff (Add a b) x = add (diff a x) (diff b x)
diff (Mul a b) x = add (mul (diff a x) b) (mul a (diff b x))


-- * Smart constructors

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 a b       = Mul a b


--------------------------------------------------------------------------------
-- | Applicative Expression Evaluator

--evalA :: Expr -> Integer
evalA (Num n)     = pure n
evalA (Add e1 e2) = (+) <$> evalA e1 <*> evalA e2
evalA (Mul e1 e2) = (*) <$> evalA e1 <*> evalA e2


--------------------------------------------------------------------------------
-- * Avoiding division by zero

evalD (Num n)     = pure n
evalD (Add e1 e2) = (+) <$> evalD e1 <*> evalD e2
evalD (Mul e1 e2) = (*) <$> evalD e1 <*> evalD e2
evalD (Div e1 e2) = do a <- evalD e1
                       b <- evalD e2
                       safeDiv1 a b

safeDiv1 :: Integer -> Integer -> Maybe Integer
safeDiv1 x 0 = Nothing
safeDiv1 x y = Just (x `div` y)


--------------------------------------------------------------------------------
-- * Making the environment implicit
{-
evalE (Num n)   = 
evalE (Var x)   = 
evalE (Add a b) = 
evalE (Mul a b) = 
--}

--------------------------------------------------------------------------------
-- * A monad for eval

-- | A monad for environments and avoiding failures

data Eval a = E ([(Name,Integer)] -> Maybe a)
runE (E f) = f

instance Functor Eval where
  fmap = liftM

instance Applicative Eval where
  pure = return
  (<*>) = ap

instance Monad Eval where
  -- return :: a -> Eval a
  return x = E ( \ env -> Just x )
  -- (>>=) :: Eval a -> (a->Eval b) -> Eval b
  E m >>= f = E ( \ env -> do x <- m env; runE (f x) env)


divByZero = E (\env-> Nothing)
lookupVar x = E (lookup x)

safeDiv x 0 = divByZero
safeDiv x y = return (x `div` y)


evalM (Num n) = pure n
evalM (Var x) = lookupVar x
evalM (Add a b) = (+) <$> evalM a <*> evalM b
evalM (Mul a b) = (*) <$> evalM a <*> evalM b
evalM (Div a b) = do x <- evalM a
                     y <- evalM b
                     safeDiv x y