module Tries where

data MapStr a = TrieStr (Maybe a) (MapChar (MapStr a))
  deriving (Eq,Show)

type MapChar a = [(Char,a)]

emptyTrieStr = TrieStr Nothing []

lookupChar :: Char -> MapChar a -> Maybe a
lookupChar = lookup

lookupStr :: String -> MapStr a -> Maybe a
lookupStr [] (TrieStr tn tcs) = tn
lookupStr (c:cs) (TrieStr tn tcs) = do
  tc <- lookupChar c tcs
  lookupStr cs tc

tidyStr :: MapStr a -> Maybe (MapStr a)
tidyStr (TrieStr Nothing []) = Nothing
tidyStr trie                 = Just trie

updateStr :: String -> (Maybe a -> Maybe a) -> MapStr a -> MapStr a
updateStr "" upd (TrieStr v tcs) = TrieStr (upd v) tcs
updateStr (c:cs) upd (TrieStr tn tc) =
  TrieStr tn
          (case lookupChar c tc of
             Nothing -> case updateStr cs upd emptyTrieStr of
                          TrieStr Nothing [] -> tc
                          tcs'                -> (c,tcs'):tc
             Just tcs -> case updateStr cs upd tcs of
                           TrieStr Nothing [] -> filter ((/=c) . fst) tc
                           tcs'               -> update c tcs' tc
          )

update :: Eq a => a -> b -> [(a,b)] -> [(a,b)]
update k v [] = [(k,v)]
update k v ((k',v'):kvs) | k == k'   = ((k,v):kvs)
			 | otherwise = (k',v'):update k v kvs

insertStr :: String -> a -> MapStr a -> MapStr a
insertStr s x = updateStr s (const (Just x))

deleteStr :: String -> MapStr a -> MapStr a
deleteStr s = updateStr s (const Nothing)

data Tree = Leaf String | Node Tree Tree
  deriving Show

data MapTree v = TrieTree (Maybe (MapStr v))
	                  (Maybe (MapTree (MapTree v)))
  deriving (Show,Eq)

lookupTree :: Tree -> MapTree a -> Maybe a
lookupTree (Leaf s) (TrieTree Nothing _) = Nothing
lookupTree (Leaf s) (TrieTree (Just tl) _) = lookupStr s tl
lookupTree (Node l r) (TrieTree _ Nothing) = Nothing
lookupTree (Node l r) (TrieTree _ (Just tn)) = do
   tn1 <- lookupTree l tn
   lookupTree r tn1

emptyTrieTree :: MapTree a
emptyTrieTree = TrieTree Nothing Nothing

tidyTree :: MapTree a -> Maybe (MapTree a)
tidyTree (TrieTree Nothing Nothing) = Nothing
tidyTree trie                       = Just trie

ensureTree :: Maybe (MapTree a) -> MapTree a
ensureTree Nothing = emptyTrieTree
ensureTree (Just trie) = trie

updateTree :: Tree -> (Maybe a -> Maybe a) -> MapTree a -> MapTree a
updateTree (Leaf s) update (TrieTree mtl tn) =
  TrieTree (tidyStr (updateStr s update (maybe emptyTrieStr id mtl)))
           tn
updateTree (Node l r) update (TrieTree tl tn) =
  TrieTree tl
           (tidyTree (updateTree l 
                       (tidyTree . updateTree r update . ensureTree)
                       (ensureTree tn)))

insertTree ::  Tree -> a -> MapTree a -> MapTree a
insertTree key v trie = updateTree key (const (Just v)) trie

deleteTree :: Tree -> MapTree a -> MapTree a
deleteTree key = updateTree key (const Nothing)

------------------------------
-- binary Numbers
------------------------------

data Bin = IHi | I Bin | O Bin
  deriving Show

data MapBin a = TrieBin (Maybe a)          -- for IHi
                        (Maybe (MapBin a)) -- for I
		        (Maybe (MapBin a)) -- for O
  deriving (Show,Eq)

emptyTrieBin :: MapBin a
emptyTrieBin = TrieBin Nothing Nothing Nothing

lookupBin :: Bin -> MapBin a -> Maybe a
lookupBin IHi (TrieBin te ti to) = te
lookupBin (I n) (TrieBin te ti to) = ti >>= lookupBin n
lookupBin (O n) (TrieBin te ti to) = to >>= lookupBin n

ensureBin :: Maybe (MapBin a) -> MapBin a
ensureBin Nothing = emptyTrieBin
ensureBin (Just trie) = trie

tidyBin ::MapBin a -> Maybe (MapBin a)
tidyBin (TrieBin Nothing Nothing Nothing) = Nothing
tidyBin trie                              = Just trie

insBin :: Bin -> a -> MapBin a -> MapBin a
insBin IHi v (TrieBin _ ti to) = TrieBin (Just v) ti to

insBin (I n) v (TrieBin te ti to) =
  TrieBin te (Just (insBin n v (ensureBin ti))) to

insBin (O n) v (TrieBin te ti to) =
  TrieBin te ti (Just (insBin n v (ensureBin to)))

delBin :: Bin -> MapBin a -> MapBin a
delBin IHi (TrieBin _ ti to) = TrieBin Nothing ti to

delBin (I n) (TrieBin te ti to) =
  maybe emptyTrieBin
        (\ti' -> TrieBin te (tidyBin (delBin n ti')) to)
        ti
delBin (O n) (TrieBin te ti to) =
  maybe emptyTrieBin
        (\to' -> TrieBin te ti (tidyBin (delBin n to')))
        to

updateBin :: Bin -> (Maybe a -> Maybe a) -> MapBin a -> MapBin a
updateBin IHi upd (TrieBin te ti to) = TrieBin (upd te) ti to
updateBin (I b) upd (TrieBin te ti to) =
  TrieBin te (tidyBin (updateBin b upd (ensureBin ti))) to
updateBin (O b) upd (TrieBin te ti to) =
  TrieBin te ti (tidyBin (updateBin b upd (ensureBin to)))

deleteBin :: Bin -> MapBin a -> MapBin a
deleteBin b = updateBin b (const Nothing)

insertBin :: Bin -> a -> MapBin a -> MapBin a
insertBin b v = updateBin b (const (Just v))

  
intToBin :: Int -> Bin
intToBin n = intToBin' n IHi
 where intToBin' 1 b = b
       intToBin' n b = intToBin' (n `div` 2)
                                 (if even n then O b else I b)