module Data.Map.TernaryMap (
                TernaryMap,
                insert,
                singleton,
                member,
                size,
                
                ) where
import Data.Bits
import Data.Binary
import Control.Monad
import Prelude hiding (lookup)


-- | TernaryMap a b is ternary tree. It is commonly used for storing word lists
-- like dictionaries.
data TernaryMap a b = Node !a !(TernaryMap a b) !(TernaryMap a b) !(TernaryMap a b)
                    | Null b !(TernaryMap a b)
                    | End
               deriving (Show, Eq)

-- | Quickly build a tree without an initial tree. This should be used
-- to create an initial tree, using insert there after.
singleton :: Ord a => [a] -> b -> TernaryMap a b
singleton (x:xs) b = Node x End (singleton xs b) End
singleton []     b = Null b End

-- | Inserts an entrie into a tree. Values with the same key will be replaced
-- with the newer value.
insert :: Ord a => [a] -> b -> TernaryMap a b -> TernaryMap a b
insert xss@(_:_)  b End              = singleton xss b
insert xss@(_:_)  b (Null b' rest)   = Null b' $ insert xss b rest
insert []         b End              = Null b End
insert []         b (Node ele l e h) = Node ele (insert [] b l) e h
insert []         b (Null _ rest)    = Null b rest
insert xss@(x:xs) b (Node ele l e h) =
    case compare x ele of
        LT -> Node ele (insert xss b l) e h
        EQ -> Node ele l (insert xs b e) h
        GT -> Node ele l e (insert xss b h)


-- | Returns true if the `[a]` is a key in the TernaryMap.
member :: Ord a => [a] -> TernaryMap a b -> Bool
member         _   End             = False
member         [] (Null _ _)       = True
member         [] (Node _ l _ _)   = member [] l
member xss@(_:_)  (Null _ rest)    = member xss rest
member xss@(x:xs) (Node ele l e h) = 
    case compare x ele of
        LT -> member xss l
        EQ -> member  xs e
        GT -> member xss h


lookup :: Ord a => [a] -> TernaryMap a b -> Maybe b
lookup _ End                       = Nothing
lookup [] (Null b _)               = Just b
lookup [] (Node _ l _ _)           = lookup [] l
lookup xs (Null _ rest)            = lookup xs rest
lookup xss@(x:xs) (Node ele l e h) =
    case compare x ele of
        LT -> lookup xss l
        EQ -> lookup  xs e
        GT -> lookup xss h

(!) :: Ord a => TernaryMap a b -> [a] -> Maybe b
(!) = flip lookup

-- | Returns the number of non-Val Elems. not exported
treeSize :: TernaryMap a b -> Int
treeSize End = 0
treeSize (Node _ l e h) = 1 + treeSize l + treeSize e + treeSize h
treeSize (Null _ rest) = treeSize rest

-- | Counts how many entries there are in the tree.
size :: TernaryMap a b -> Int
size End            = 0
size (Node _ l e h) = size l + size e + size h
size (Null _ rest)  = 1 + size rest

-- | Creates a new tree from a list of 'strings'
fromList :: Ord a => [([a],b)] -> TernaryMap a b
fromList = foldl (\tree (as,b) -> insert as b tree) empty

-- | An empty map.
empty :: TernaryMap a b
empty = End

-- | Makes a list of all the values in the map.
elems :: TernaryMap a b -> [b]
elems  End           = []
elems (Node _ l e h) = elems l ++ (elems e ++ elems h)
elems (Null b rest)  = b : elems rest

-- | Returns true if the map is empty.
null :: TernaryMap a b -> Bool
null End = True
null _   = False

-- keySet :: TernaryMap a b -> S.TernarySet a
-- keySet End = S.End
-- keySet (Node (C x) l e h) = S.Node (S.C x) (keySet l) (keySet e) (keySet h)
-- keySet (Node (Val _) l e h) = S.Node (S.Null) (keySet l) (keySet e) (keySet h)


instance Functor (TernaryMap a) where
    fmap _ End = End
    fmap f (Null b rest)    = Null (f b) (fmap f rest)
    fmap f (Node ele l e h) = Node ele (fmap f l) (fmap f e) (fmap f h)

-- | A rather long Binary instance, that uses binary numbers to indicate
-- where Ends are efficiently.
instance (Binary a, Binary b) => Binary (TernaryMap a b) where
    put (Node ch End End End) = do
        putWord8 0
        put ch
    put (Node ch End End h) = do
        putWord8 1
        put ch
        put h
    put (Node ch End e End) = do
        putWord8 2
        put ch
        put e
    put (Node ch End e h) = do
        putWord8 3
        put ch
        put e
        put h
    put (Node ch l End End) = do
        putWord8 4
        put ch
        put l
    put (Node ch l End h) = do
        putWord8 5
        put ch
        put l
        put h
    put (Node ch l e End) = do
        putWord8 6
        put ch
        put l
        put e
    -- General case
    put (Node ch l e h) = do
        putWord8 7
        put ch
        put l
        put e
        put h
    put (Null b End) = putWord8 8 >> put b
    put (Null b rest) = do
        putWord8 9
        put b
        put rest
    put End = putWord8 10

    get = do
        tag <- getWord8
        case tag of
            _ | tag < 8 ->
                do
                    ch <- get
                    l <- if (tag `testBit` 2) then get else return End
                    e <- if (tag `testBit` 1) then get else return End
                    h <- if (tag `testBit` 0) then get else return End
                    return (Node ch l e h)
            8 -> liftM (flip Null End) get
            9 -> liftM2 Null get get
            _ -> return End