import Test.QuickCheck
import Control.Monad

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

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

-- try:
-- > quickCheck prop_basic1

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

-- try:
-- > quickCheck prop_basic2

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

-- try:
-- > quickCheck prop_basic3

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

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

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

checkout :: [Pack] -> Double
checkout = sum . map price
  where price EmptyPack = 0
        price (Tomatoes weight) = weight * 33.50
        price (Cucumbers count) =
          fromIntegral count * 19.90

instance Arbitrary Pack where
  arbitrary = packGen1

packGen1 :: Gen Pack
packGen1 = oneof [ pure EmptyPack
                 , fmap Tomatoes arbitrary
                 , fmap Cucumbers arbitrary ]

prop_pack1 :: [Pack] -> Bool
prop_pack1 pack = checkout pack >= 0

-- try:
-- > quickCheck prop_pack1
-- > sample packGen1

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

-- try:
-- > quickCheck prop_pack2
-- > verboseCheck prop_pack2

packGen2 :: Gen Pack
packGen2 = oneof
  [ pure EmptyPack
  , fmap Tomatoes (arbitrary `suchThat` (>=0))
  , fmap Cucumbers (arbitrary `suchThat` (>=0)) ]

-- try:
-- > sample packGen2
-- > quickCheck $ forAll (listOf packGen1) prop_pack2
-- > quickCheck $ forAll (listOf packGen2) prop_pack2
-- > quickCheck $ forAll (listOf packGen2) prop_pack1

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

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

instance Arbitrary BinTree where
  arbitrary = treeGen1
  shrink = treeShrink

treeGen1 :: Gen BinTree
treeGen1 = oneof [ pure BEmpty,
         liftM3 BNode arbitrary treeGen1 treeGen1 ]

treeShrink :: BinTree -> [BinTree]
treeShrink BEmpty = []
treeShrink (BNode v l r) = [l, r]
                ++ [ BNode v' l r | v' <- shrink v ]
                ++ [ BNode v l' r | l' <- shrink l ]
                ++ [ BNode v l r' | r' <- shrink r ]

-- try:
-- > sample 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

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

-- try:
-- > quickCheck prop_tree1
-- > verboseCheck prop_tree1

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

-- try:
-- > quickCheck prop_tree2

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

-- try:
-- > quickCheck prop_tree3

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

-- try:
-- > quickCheck prop_tree4

treeGen2 :: Gen BinTree
treeGen2 = frequency
  [ (1, pure BEmpty)
  , (4, liftM3 BNode arbitrary treeGen2 treeGen2) ]

-- try:
-- > sample treeGen2

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

treeGen3 :: Gen BinTree
treeGen3 = sized treeGen where
  treeGen 0 = pure BEmpty
  treeGen n = frequency
    [ (1, pure BEmpty)
    , (4, liftM3 BNode arbitrary subtree subtree) ]
       where subtree = treeGen (n `div` 2)

-- try:
-- > sample treeGen3
-- > quickCheck $ forAll treeGen3 prop_tree4

treeGen4 :: Gen BinTree
treeGen4 = sized treeGen
    where treeGen 0 = pure BEmpty
          treeGen n = frequency [ (1, pure BEmpty)
                                , (n, liftM3 BNode arbitrary subtree subtree) ]
                      where subtree = treeGen (n `div` 2)

-- try:
-- > sample treeGen4
-- > quickCheck $ forAll treeGen4 prop_tree4