import Test.QuickCheck
import Control.Monad

-- ==== basic example ====

-- try: 
-- > quickCheck prop_basic1
prop_basic1 :: Int -> [a] -> Bool
prop_basic1 n xs = length (take n xs) == n  

-- try:
-- > quickCheck prop_basic2
prop_basic2 :: Int -> [a] -> Bool
prop_basic2 n xs = length (take n xs) <= n

newtype NonNegativeInt = NonNegativeInt Int
                         deriving Show

-- try:
-- > sample nonNegativeGen
nonNegativeGen :: Gen NonNegativeInt
nonNegativeGen = fmap NonNegativeInt $ 
                   choose (0, 50 :: Int)

instance Arbitrary NonNegativeInt where
  arbitrary = nonNegativeGen

-- try:
-- > quickCheck prop_basic3
prop_basic3 :: NonNegativeInt -> [a] -> Bool
prop_basic3 (NonNegativeInt n) xs = length (take n xs) <= n

-- try:
-- > quickCheck prop_basic4
prop_basic4 :: NonNegative Int -> [a] -> Bool
prop_basic4 (NonNegative n) xs = length (take n xs) <= n

-- try:
-- > quickCheckResult prop_basic4
-- > quickCheckResult prop_basic2
-- > quickCheckWith (stdArgs { maxSuccess = 200 }) prop_basic4
-- > verboseCheck prop_basic4

-- ==== Shopping example ====

data Pack = EmptyPack        -- empty pack
          | Tomatoes Double  -- tomato weight in kg
          | Cucumbers Int    -- number of cucumbers
          deriving (Eq, Show)

instance Arbitrary Pack where
  arbitrary = packGen1

-- try:
-- > sample packGen1
packGen1 :: Gen Pack
packGen1 = oneof [ return EmptyPack
                 , fmap Tomatoes arbitrary
                 , fmap Cucumbers arbitrary ]

price1 :: Pack -> Double
price1 EmptyPack = 0
price1 (Tomatoes weight) = weight * 33.50
price1 (Cucumbers count) = fromIntegral count * 19.90

checkout1 :: [Pack] -> Double
checkout1 = sum . map price1

checkout2 [] = 0
checkout2 (EmptyPack : xs) = checkout2 xs
checkout2 (Tomatoes w : xs) = w * 33.5 + checkout2 xs
checkout2 (Cucumbers n : xs) = fromIntegral n * 19.9 + checkout2 xs

-- try:
-- > quickCheck prop_pack1
prop_pack1 :: [Pack] -> Bool
prop_pack1 pack = checkout1 pack >= 0

-- try:
-- > quickCheck prop_pack2
-- > verboseCheck prop_pack2
prop_pack2 :: [Pack] -> Property
prop_pack2 pack = all nonNegative pack ==> checkout1 pack >= 0
  where nonNegative (Tomatoes w) = w >= 0
        nonNegative (Cucumbers n) = n >= 0
        nonNegative _ = True

-- try:
-- > sample packGen2
-- > quickCheck $ forAll (listOf packGen1) prop_pack2
-- > quickCheck $ forAll (listOf packGen2) prop_pack2
-- > quickCheck $ forAll (listOf packGen2) prop_pack1
packGen2 :: Gen Pack
packGen2 = oneof 
  [ return EmptyPack
  , fmap Tomatoes (arbitrary `suchThat` (>=0))
  , fmap Cucumbers (arbitrary `suchThat` (>=0)) ]

-- try: 
-- > quickCheck prop_pack3
prop_pack3 :: [Pack] -> Property
prop_pack3 pack = checkout1 pack === checkout2 pack

-- try:
-- > quickCheck prop_pack4
prop_pack4 :: [Pack] -> Property
prop_pack4 pack = round (checkout1 pack) === round (checkout2 pack)

-- ==== Tree example ====

data BinTree = BEmpty
             | BNode Int BinTree BinTree
             deriving (Eq, Ord, Show)

instance Arbitrary BinTree where
  arbitrary = treeGen1

-- try:
-- > sample treeGen1
treeGen1 :: Gen BinTree
treeGen1 = oneof 
  [ return BEmpty
  , liftM3 BNode arbitrary treeGen1 treeGen1 ]

treeToList :: BinTree -> [Int]
treeToList BEmpty = []
treeToList (BNode v l r) = v : treeToList l ++ treeToList r

treeSum :: BinTree -> Int
treeSum BEmpty = 0
treeSum (BNode v l r) = v + treeSum l + treeSum r

-- try:
-- > quickCheck prop_tree1
-- > verboseCheck prop_tree1
prop_tree1 :: BinTree -> Bool
prop_tree1 t = treeSum t == sum (treeToList t)

-- try:
-- > quickCheck prop_tree2
prop_tree2 :: BinTree -> Property
prop_tree2 t = classify (treeSize t == 0) "trivial" $ prop_tree1 t

-- try:
-- > quickCheck prop_tree3
prop_tree3 :: BinTree -> Property
prop_tree3 t = classify (treeSize t == 1) "easy" $ prop_tree2 t

-- try:
-- > quickCheck prop_tree4
prop_tree4 :: BinTree -> Property
prop_tree4 t = collect (treeSize t) $ prop_tree3 t

-- try:
-- > sample treeGen2
treeGen2 :: Gen BinTree
treeGen2 = frequency 
  [ (1, return BEmpty)
  , (4, liftM3 BNode arbitrary treeGen2 treeGen2) ]

treeSize :: BinTree -> Int
treeSize BEmpty = 0
treeSize (BNode _ l r) = 1 + treeSize l + treeSize r

-- try:
-- > sample treeGen3
-- > quickCheck $ forAll treeGen3 prop_tree4
treeGen3 :: Gen BinTree
treeGen3 = sized treeGen where 
  treeGen 0 = return BEmpty
  treeGen n = frequency 
    [ (1, return BEmpty)
    , (4, liftM3 BNode arbitrary subtree subtree) ]
       where subtree = treeGen (n `div` 2)

-- try:
-- > sample treeGen4
-- > quickCheck $ forAll treeGen4 prop_tree4
treeGen4 :: Gen BinTree
treeGen4 = sized treeGen
    where treeGen 0 = return BEmpty
          treeGen n = frequency [ (1, return BEmpty)
                                , (n, liftM3 BNode arbitrary subtree subtree) ]
                      where subtree = treeGen (n `div` 2)