import Test.QuickCheck
import Data.List
import Data.Maybe

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

data Expr
  = Num Integer
  | Add Expr Expr
  | Mul Expr Expr
  | Var Name
 deriving ( Eq )

type Name = String

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

showExpr :: Expr -> String
showExpr (Num n)   = show n
showExpr (Add a b) = showExpr a ++ "+" ++ showExpr b
showExpr (Mul a b) = showFactor a ++ "*" ++ showFactor b
showExpr (Var x)   = x

showFactor :: Expr -> String
showFactor (Add a b) = "(" ++ showExpr (Add a b) ++")"
showFactor e         = showExpr e

instance Show Expr where
  show = showExpr

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

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

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

instance Arbitrary Expr where
  arbitrary = sized arbExpr

arbExpr :: Int -> Gen Expr
arbExpr s =
  frequency [ (1, do n <- arbitrary
                     return (Num n))
            , (s, do a <- arbExpr s'
                     b <- arbExpr s'
                     return (Add a b))
            , (s, do a <- arbExpr s'
                     b <- arbExpr s'
                     return (Mul a b))
            , (1, do x <- elements ["x","y","z"]
                     return (Var x))
            ]
 where
  s' = s `div` 2

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

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

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

diff :: Expr -> Name -> Expr
diff (Num n)   x = 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 b (diff a x))
diff (Var y)   x
  | x == y       = Num 1
  | otherwise    = Num 0

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

simplify :: Expr -> Expr
simplify e | null (vars e) =
  Num (eval [] e)

simplify e = e

isNum :: Expr -> Bool
isNum (Num _) = True
isNum _       = False

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

data Env = Env [(Name,Integer)]
 deriving ( Show )

instance Arbitrary Env where
  arbitrary =
    do x <- arbitrary
       y <- arbitrary
       z <- arbitrary
       return (Env [("x",x),("y",y),("z",z)])

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

prop_SimplifyCorrect e (Env env) =
  eval env e == eval env (simplify e)

prop_SimplifyNoJunk e =
  noJunk (simplify e)
 where
  -- just an example, can be extended at will!
  noJunk (Add a b) = not (isNum a && isNum b)
                  && noJunk a && noJunk b
  noJunk (Mul a b) = not (isNum a && isNum b)
                  && noJunk a && noJunk b
  noJunk _         = True

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