{-# LANGUAGE ScopedTypeVariables, TemplateHaskell #-}

module Main where

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

-- imports

import Test.QuickCheck

import Text.Show.Functions

import Data.List

  ( sort

  , group

  , nub

  , (\\)

  )

import Control.Monad

  ( liftM

  , liftM2

  )

import Data.Maybe

--import Text.Show.Functions

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

-- binary search trees

data Set a

  = Node a (Set a) (Set a)

  | Empty

 deriving ( Eq, Ord, Show )

empty :: Set a

empty = Empty

isEmpty :: Set a -> Bool

isEmpty Empty = True

isEmpty _     = False

unit :: a -> Set a

unit x = Node x empty empty

size :: Set a -> Int

size Empty          = 0

size (Node _ s1 s2) = 1 + size s1 + size s2

insert :: Ord a => a -> Set a -> Set a

insert x s = s `union` unit x

merge :: Set a -> Set a -> Set a

s `merge` Empty                      = s

s `merge` Node x Empty s2            = Node x s s2

s `merge` Node x (Node y s11 s12) s2 = Node y s (Node x (s11 `merge` s12) s2)

delete :: Ord a => a -> Set a -> Set a

delete x Empty = Empty

delete x (Node x' s1 s2) =

  case x `compare` x' of

    LT -> Node x' (delete x s1) s2

    EQ -> s1 `merge` s2

    GT -> Node x' s1 (delete x s2)

union :: Ord a => Set a -> Set a -> Set a

{-

s1    `union` Empty = s1

Empty `union` s2    = s2

s1@(Node x s11 s12) `union` s2@(Node y s21 s22) =

  case x `compare` y of

    LT -> Node x s11 (s12 `union` Node y Empty s22) `union` s21

    EQ -> Node x (s11 `union` s21) (s12 `union` s22)

    --GT -> s11 `union` Node y s21 (Node x Empty s12 `union` s22)

    GT -> Node x (s11 `union` Node y s21 Empty) s12 `union` s22 

-}

s1             `union` Empty = s1

Empty          `union` s2    = s2

Node x s11 s12 `union` s2    = Node x (s11 `union` s21) (s12 `union` s22)

 where

  (s21,s22) = split x s2

split :: Ord a => a -> Set a -> (Set a, Set a)

split x Empty = (Empty, Empty)

split x (Node y s1 s2) =

  case x `compare` y of

    LT -> (s11, Node y s12 s2)

    EQ -> (s1, s2)

    GT -> (Node y s1 s21, s22)

 where

  (s11,s12) = split x s1

  (s21,s22) = split x s2

mapp :: (a -> b) -> Set a -> Set b

mapp f Empty          = Empty

mapp f (Node x s1 s2) = Node (f x) (mapp f s1) (mapp f s2)

fromList :: Ord a => [a] -> Set a

--fromList xs = build [ (empty,x) | x <- sort xs ]

fromList xs = build [ (empty,head x) | x <- group (sort xs) ]

 where

  build []      = empty

  build [(s,x)] = attach x s

  build sxs     = build (sweep sxs)

  sweep []                    = []

  sweep [sx]                  = [sx]

  sweep ((s1,x1):(s2,x2):sxs) = (Node x1 s1 s2,x2) : sweep sxs

  attach x Empty          = unit x

  attach x (Node y s1 s2) = Node y s1 (attach x s2)

toList :: Set a -> [a]

toList s = toSortedList s

toSortedList :: Set a -> [a]

toSortedList s = toList' s []

 where

  toList' Empty          xs = xs

  toList' (Node x s1 s2) xs = toList' s1 (x : toList' s2 xs)

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

-- generators

instance (Ord a, Arbitrary a) => Arbitrary (Set a) where

  arbitrary = sized (arbSet Nothing Nothing)

   where

    arbSet mx my n =

      frequency $

        [ (1, return Empty) ] ++

        [ (7, do mz <- arbitrary `suchThatMaybe` (isOK mx my)

                 case mz of

                   Nothing -> return Empty

                   Just z  -> liftM2 (Node z) (arbSet mx mz n2)

                                              (arbSet mz my n2)

                    where n2 = n `div` 2)

        | n > 0

        ]

    isOK mx my z =

      maybe True (

  shrink Empty            = []

  shrink t@(Node x s1 s2) = [ s1, s2 ]

                         ++ [ t' | x' <- shrink x, let t' = Node x' s1 s2, invariant t' ]

-- instance (Ord a, ShrinkSub a) => ShrinkSub (Set a)

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

-- properties

(.<) :: Ord a => Set a -> a -> Bool

Empty      .< x = True

Node y _ s .< x = y < x && s .< x

(<.) :: Ord a => a -> Set a -> Bool

x <. Empty      = True

x <. Node y _ s = x < y && x <. s

(==?) :: Ord a => Set a -> [a] -> Bool

s ==? xs = invariant s && sort (toList s) == nub (sort xs)

invariant :: Ord a => Set a -> Bool

invariant Empty          = True

invariant (Node x s1 s2) = s1 .< x && x <. s2 && invariant s1 && invariant s2

prop_Invariant (s :: Set Int) =

  invariant s

prop_Empty =

  empty ==? ([] :: [Int])

prop_Unit (x :: Int) =

  unit x ==? [x]

prop_Size (s :: Set Int) =

  cover (size s >= 15) 60 "large" $

    size s == length (toList s)

prop_Insert x (s :: Set Int) =

  insert x s ==? (x : toList s)

prop_Delete x (s :: Set Int) =

  delete x s ==? (toList s \\ [x])

prop_Union s1 (s2 :: Set Int) =

  (s1 `union` s2) ==? (toList s1 ++ toList s2)

prop_Mapp (f :: Int -> Int) (s :: Set Int) =

  expectFailure $

    whenFail (putStrLn ("Fun: " ++ show [ (x,f x) | x <- toList s])) $

      mapp f s ==? map f (toList s)

prop_FromList (xs :: [Int]) =

  fromList xs ==? xs

prop_ToSortedList (s :: Set Int) =

  s ==? xs && xs == sort xs

 where

  xs = toSortedList s

--  whenFail (putStrLn ("Result: " ++ show (fromList xs))) $

prop_FromList' (xs :: [Int]) =

  shrinking shrink xs $ \xs' ->

    fromList xs ==? xs

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

-- main

return []

main = $quickCheckAll

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

-- the end.