-- Implements a Map ADT using a Binary Search Tree data structure
-- by Scot Drysdale

module BSTMap (Map,   -- Note we don't export constructors
              (!), lookup, insert, member, insertWith, delete, size, empty, 
              keys, toList, fromList
            ) where

import Prelude hiding (lookup)  -- We re-define lookup here

data Map k a = Leaf 
             | Branch {key :: k, value :: a, 
                       left :: (Map k a), right :: (Map k a)}
     deriving Show

-- Create an emptyMap
empty :: Map a b
empty = Leaf

-- Insert key k and value v into this Map.  
-- Replace value with (f new_val old_val) if key already present.
insertWith ::  Ord a => (b -> b -> b) -> a -> b -> Map a b -> Map a b
insertWith f k v Leaf = Branch k v Leaf Leaf
insertWith f k v (Branch k1 v1 l r) 
           | k == k1 = Branch k (f v v1) l r  -- Replace value using f
           | k > k1  = Branch k1 v1 l (insertWith f k v r)
           | k < k1  = Branch k1 v1 (insertWith f k v l) r

-- Insert key k and value v into thisMap.  
-- Replace value with new value if key already present.
insert         :: Ord a => a -> b -> Map a b -> Map a b
insert k v tree = insertWith (\x y->x) k v tree

-- Find the key and return the value, or Nothing if not present 
lookup :: Ord a => a -> Map a b -> Maybe b
lookup k Leaf                        = Nothing
lookup k (Branch k1 v l r) | k == k1 = Just v
                           | k > k1  = lookup k r
                           | k < k1  = lookup k l
                         
(!) :: Ord a => Map a b -> a -> b
t ! k = case lookup k t of
          Nothing -> error "Key not found in Map"
          (Just v) -> v
          
-- Determines if key is in the Map
member    :: Ord a => a -> Map a b -> Bool
member k t = case lookup k t of
               Nothing -> False
               (Just _) -> True

-- Delete key from Map.                          
delete :: Ord a => a -> Map a b -> Map a b
delete k Leaf = Leaf
delete k (Branch k1 v l r) 
         | k > k1 = Branch k1 v l (delete k r)
         | k < k1 = Branch k1 v (delete k l) r
         | k == k1 = case (l, r) of
             (Leaf, _) -> r
             (_, Leaf) -> l
             (_, _)    -> 
                let (k2, v2) = maxPair l
                in Branch k2 v2 (delete k2 l) r

-- Finds the pair with the maximum key.
-- Only works on non-empty trees. 
maxPair :: Map a b -> (a, b)
maxPair (Branch k v _ Leaf) = (k, v)
maxPair (Branch _ _ _ r ) = maxPair r

-- Converts a tree into a list of (key, value) pairs
toList :: Map a b -> [(a,b)]
toList tree = toListHelper tree []
  where toListHelper :: Map a b -> [(a,b)] -> [(a,b)]
        toListHelper Leaf acc = acc
        toListHelper (Branch k v l r) acc = 
          toListHelper l ((k,v) : toListHelper r acc)
          
-- Converts a list of (key, value) pairs to a Map
fromList  :: (Ord a) => [(a,b)] -> Map a b
fromList = foldr (uncurry insert) empty

-- Returns a list of the keys in the Map
keys  :: Ord a => Map a b -> [a]
keys = map fst . toList

-- Returns the number of keys in theMap
size                 :: Map a b -> Int
size Leaf             = 0
size (Branch _ _ l r) = 1 + size l + size r

-- Prints the tree, using indentation to show depth
showTree :: (Show a, Show b) => Map a b -> IO ()
showTree t = prettyPrint t ""

-- Helper function for showTree
prettyPrint :: (Show a, Show b) => Map a b -> String -> IO ()
prettyPrint Leaf indent = return ()
prettyPrint (Branch k v l r) indent = do
       prettyPrint r (indent ++ "   ") 
       putStr (indent ++ (show k) ++ "  " ++ (show v) ++ "\n")
       prettyPrint l (indent ++ "   ")
       

-- a map of things to their weigh in pounds       
t =  insert "boar" 500 (insert "car" 3000 (insert "cow" 890 (insert "bee" 0.01 
            (insert "dog" 20 (insert "cat" 8.0 (Leaf))))))