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

--------------------------------------------------------------------------
-- 0. Exercises from the book

size :: Expr -> Int
size (Lit _)   = 0
size (Add a b) = 1 + size a + size b
size (Sub a b) = 1 + size a + size b

---

data Expr
  = Lit Int
  | Add Expr Expr
  | Sub Expr Expr
  | Mul Expr Expr
  | Div Expr Expr
 deriving ( Eq, Show )

eval :: Expr -> Int
eval (Lit n)   = n
eval (Add a b) = eval a + eval b
eval (Sub a b) = eval a - eval b
eval (Mul a b) = eval a * eval b
eval (Div a b) = eval a `div` eval b

---

evalMaybe :: Expr -> Maybe Int
evalMaybe (Lit n)   = Just n

evalMaybe (Add a b) = evalMaybe a `add` evalMaybe b
 where
  Nothing `add` _       = Nothing
  _       `add` Nothing = Nothing
  Just x  `add` Just y  = Just (x+y)

evalMaybe (Sub a b) = evalMaybe a `sub` evalMaybe b
 where
  Nothing `sub` _       = Nothing
  _       `sub` Nothing = Nothing
  Just x  `sub` Just y  = Just (x-y)
  
evalMaybe (Mul a b) = evalMaybe a `mul` evalMaybe b
 where
  Nothing `mul` _       = Nothing
  _       `mul` Nothing = Nothing
  Just x  `mul` Just y  = Just (x*y)
  
evalMaybe (Div a b) = evalMaybe a `dvv` evalMaybe b
 where
  Nothing `dvv` _        = Nothing
  _       `dvv` Nothing  = Nothing
  _       `dvv` (Just 0) = Nothing
  Just x  `dvv` Just y   = Just (x `div` y)
  
---

data Expr'
  = Lit' Int
  | Op Ops Expr' Expr'
 deriving ( Eq, Show )
  
data Ops = Add' | Sub' | Mul' | Div' | Mod'
 deriving ( Eq, Show )

eval' :: Expr' -> Maybe Int
eval' (Lit' n)    = Just n
eval' (Op op a b) = ops op (eval' a) (eval' b)
 where
  ops _   Nothing  _                   = Nothing
  ops _   _        Nothing             = Nothing
  ops op  _        (Just 0) | isDiv op = Nothing
  ops op  (Just x) (Just y)            = Just (evalOp op x y)
  
  isDiv op = op `elem` [Div',Mod']
  
  evalOp Add' x y = x + y
  evalOp Sub' x y = x - y
  evalOp Mul' x y = x * y
  evalOp Div' x y = x `div` y
  evalOp Mod' x y = x `mod` y

---

data Exp
  = Lt Int
  | Exp :+: Exp
  | Exp :-: Exp
 deriving ( Eq, Show )

siz :: Exp -> Int
siz (Lt _)    = 0
siz (a :+: b) = 1 + siz a + siz b
siz (a :-: b) = 1 + siz a + siz b

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

data NTree
  = NilT
  | Node Int NTree NTree
 deriving ( Eq, Show )

---

left, right :: NTree -> NTree
left  (Node _ l _) = l
right (Node _ _ r) = r

---

sumTree :: NTree -> Int
sumTree NilT         = 0
sumTree (Node x l r) = x + sumTree l + sumTree r

---

sizeTree :: NTree -> Int
sizeTree NilT         = 0
sizeTree (Node x l r) = 1 + sizeTree l + sizeTree r

---

collapse :: NTree -> [Int]
collapse NilT         = []
collapse (Node x l r) = collapse l ++ [x] ++ collapse r

---

maxT, minT :: NTree -> Int
maxT t = maximum (collapse t)
minT t = minimum (collapse t)

---

reflect :: NTree -> NTree
reflect NilT         = NilT
reflect (Node x l r) = Node x (reflect r) (reflect l)

---

sort' :: NTree -> [Int]
sort' t = Data.List.sort (collapse t)

---

prop_ReflectReflect t =
  reflect (reflect t) == t

prop_CollapseLength t =
  length (collapse t) == sizeTree t

prop_CollapseSum t =
  sum (collapse t) == sumTree t

prop_ReflectCollapse t =
  collapse (reflect t) == reverse (collapse t)

---

instance Arbitrary NTree where
  arbitrary = sized arbNTree

arbNTree :: Int -> Gen NTree
arbNTree s =
  frequency [ (1, do return NilT)
            , (s, do n <- arbitrary
                     l <- arbNTree s'
                     r <- arbNTree s'
                     return (Node n l r))
            ]
 where
  s' = s `div` 2

---------------------------------------------------------------------
-- 1. Tables

build :: Ord a => [(a,b)] -> Table a b
build xys = table (length xys') xys'
 where
  xys' = sortBy first' xys                  -- bug!
  --xys' = nubBy first (sortBy first' xys)  -- correct

  (x1,_) `first`  (x2,_) = x1 == x2
  (x1,_) `first'` (x2,_) = x1 `compare` x2

  table _ []  = Empty
  table n xys = Join (table k lefts) x y (table (n-k-1) rights)
   where
    k      = n `div` 2
    lefts  = take k xys
    rights = drop (k+1) xys
    (x,y)  = xys !! k

prop_Build_Invariant xys =
  invariant (build (xys :: [(Int,Double)]))
    
prop_Build_Complete xys =
  and [ (x,y) `elem` (xys :: [(Int,Double)]) | (x,y) <- contents (build xys) ]
    
prop_Build_Enough xys =
  and [ x `elem` map fst (contents (build xys)) | (x,_) <- xys :: [(Int,Double)] ]
    
invariant :: Ord a => Table a b -> Bool
invariant Empty          = True
invariant (Join l x y r) =
  l <. x && x .< r && invariant l && invariant r
 where
  Empty         <. x = True
  Join _ x' _ _ <. x = x' < x

  x .< Empty         = True
  x .< Join _ x' _ _ = x < x'

--

data Table a b
  = Empty
  | Join (Table a b) a b (Table a b)
 deriving ( Eq, Show )

contents :: Table a b -> [(a,b)]
contents Empty          = []
contents (Join l x y r) = contents l ++ [(x,y)] ++ contents r

--------------------------------------------------------------------------
-- 2. File Systems

data File
  = File String
  | Dir String [File]
 deriving ( Eq, Show )

type FileSystem = [File]

-- this function returns all paths
search :: FileSystem -> String -> [String]
search files name =
  [ name
  | File name' <- files
  , name == name'
  ] ++
  [ dir ++ "/" ++ path
  | Dir dir files' <- files
  , path <- search files' name
  ]

-- this function returns maybe a path
searchMaybe :: FileSystem -> String -> Maybe String
searchMaybe files name =
  listToMaybe
  ( [ name
    | File name' <- files
    , name == name'
    ] ++
    [ dir ++ "/" ++ path
    | Dir dir files' <- files
    , Just path <- [searchMaybe files' name]
    ]
  )

-- it can also be defined using the first one...
searchMaybe' :: FileSystem -> String -> Maybe String
searchMaybe' files name = listToMaybe (search files name)

exampleFileSystem :: FileSystem
exampleFileSystem =
  [ File "apa"
  , Dir "bepa" [ File "apa", Dir "bepa" [], Dir "cepa" [ File "bepa" ] ]
  , Dir "cepa" [ Dir "bepa" [], Dir "cepa" [ File "apa" ] ]
  ]



---------------------------------------------------------------------
-- 3. Simple Sets

data Set a = Set [a]
 deriving ( Show )
 
empty :: Set a
empty = Set []

add :: Eq a => a -> Set a -> Set a
add x (Set xs) | x `elem` xs = Set xs
               | otherwise   = Set (x:xs)

remove :: Eq a => a -> Set a -> Set a
remove x (Set xs) = Set (xs \\ [x])

combine :: Eq a => Set a -> Set a -> Set a
Set xs `combine` Set ys = Set (xs ++ ys)   -- bug!
--Set xs `combine` Set ys = Set (xs `union` ys) -- correct

member :: Eq a => a -> Set a -> Bool
member x (Set xs) = x `elem` xs

isEmpty :: Set a -> Bool
isEmpty (Set xs) = null xs

--

instance (Arbitrary a, Eq a) => Arbitrary (Set a) where
  arbitrary =
    do xs <- arbitrary
       return (Set (nub xs))

-- set invariant      
inv :: Set Int -> Bool
inv (Set xs) = nub xs == xs

-- set contents     
conts :: Set Int -> [Int]
conts (Set xs) = xs

-- helper function:
-- xs =*= ys checks if the lists xs and ys contain the same elements,
-- disregarding the number of times these elements occur
(=*=) :: [Int] -> [Int] -> Bool
xs =*= ys = sort (nub xs) == sort (nub ys)

prop_Empty =
  conts empty =*= []
  
prop_Add x s =
  conts (add x s) =*= (x : conts s)
  
prop_Remove x s =
  conts (remove x s) =*= (conts s \\ [x])
  
prop_Combine s1 s2 =
  conts (s1 `combine` s2) =*= (conts s1 ++ conts s2)
  
prop_Member x s =
  member x s == elem x (conts s)
  
prop_IsEmpty s =
  isEmpty s == null (conts s)

--

prop_Inv s =
  inv s

prop_Empty_Inv =
  inv empty
  
prop_Add_Inv x s =
  inv (add x s)
  
prop_Remove_Inv x s =
  inv (remove x s)  

prop_Combine_Inv s1 s2 =
  inv (s1 `combine` s2)  

---------------------------------------------------------------------
-- 4. Sorted Sets

{-
data Set a = Set [a]
 deriving ( Show )
 
empty :: Set a
empty = Set []

add :: Ord a => a -> Set a -> Set a
add x (Set xs) | x `elem` xs = Set xs
               | otherwise   = Set (insert x xs)

remove :: Eq a => a -> Set a -> Set a
remove x (Set xs) = Set (xs \\ [x])

combine :: Ord a => Set a -> Set a -> Set a
Set xs `combine` Set ys = Set (xs `merge` ys)
 where
  [] `merge` ys = ys
  xs `merge` [] = xs
  (x:xs) `merge` (y:ys)
    | x < y     = x : (xs `merge` (y:ys))
    | x == y    = x : (xs `merge` ys)
    | otherwise = y : ((x:xs) `merge` ys)

member :: Eq a => a -> Set a -> Bool
member x (Set xs) = x `elem` xs

isEmpty :: Set a -> Bool
isEmpty (Set xs) = null xs

--

instance (Arbitrary a, Ord a) => Arbitrary (Set a) where
  arbitrary =
    do xs <- arbitrary
       return (Set (sort (nub xs)))

-- set invariant      
inv :: Set Int -> Bool
inv (Set xs) = nub xs == xs && sort xs == xs

-- set contents     
conts :: Set Int -> [Int]
conts (Set xs) = xs

-- the rest stays the same!
-}

---------------------------------------------------------------------
-- 5. Tree Sets

{-
data Set a
  = Nil
  | Comb (Set a) a (Set a)
 deriving ( Show )

empty :: Set a
empty = Nil

add :: Ord a => a -> Set a -> Set a
add x Nil     = Comb Nil x Nil
add x (Comb l y r)
  | x < y     = Comb (add x l) y r
  | x == y    = Comb l y r
  | otherwise = Comb l y (add x r)

remove :: Ord  a => a -> Set a -> Set a
remove x Nil  = Nil
remove x (Comb l y r)
  | x < y     = Comb (remove x l) y r
  | x == y    = l `combine` r
  | otherwise = Comb l y (remove x r)

-- there are many choices here... this is a very simple
-- and not very efficient combine function
combine :: Ord a => Set a -> Set a -> Set a
s1 `combine` s2 = comb [s1,s2]
 where
  comb []                = Nil
  comb (Nil : ss)        = comb ss
  comb (Comb l x r : ss) = add x (comb (l:r:ss))

member :: Ord a => a -> Set a -> Bool
member x Nil  = False
member x (Comb l y r)
  | x < y     = member x l
  | x == y    = True
  | otherwise = member x r

isEmpty :: Set a -> Bool
isEmpty Nil = True
isEmpty _   = False

--

instance (Arbitrary a, Ord a) => Arbitrary (Set a) where
  arbitrary =
    do xs <- arbitrary
       return (set xs)
   where
    set [] = Nil
    set (x:xs) = Comb (set smaller) x (set larger)
     where
      smaller = [y | y <- xs, y < x]
      larger  = [y | y <- xs, y > x]

-- set invariant      
inv :: Set Int -> Bool
inv Nil          = True
inv (Comb l x r) = l <. x && x .< r && inv l && inv r
 where
  Nil        <. x = True
  Comb l y r <. x = y < x
  
  x .< Nil        = True
  x .< Comb l y r = x < y
  
-- set contents     
conts :: Set Int -> [Int]
conts Nil          = []
conts (Comb l x r) = conts l ++ [x] ++ conts r

-- the rest stays the same!
-}

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