module Data.Map.TernaryMap (
                TernaryMap,
                insert,
                singleton,
                member,
                size,
                
                ) where
import Data.Binary
import Control.Monad
import qualified Data.Set.TernarySet as S
import Prelude hiding (lookup)


-- | Elem2 a b is used to hold elements of a list after insertion, and
-- indicate that we've reached the end of the list.
data Elem2 a b = C !a
               | Val b
             deriving (Show)
-- | TernaryMap a b is ternary tree. It is commonly used for storing word lists
-- like dictionaries.
data TernaryMap a b = TNode !(Elem2 a b) !(TernaryMap a b) !(TernaryMap a b) !(TernaryMap a b)
                    | TEnd
               deriving (Show, Eq)


instance Eq a => Eq (Elem2 a b) where
    (Val _) == (Val _) = True
    (Val _) == x        = False
    x        == (Val _) = False
    (C a)    == (C b)    = a == b

-- | All elements are greater than the Val Elem, otherwise they are
-- ordered according to their own ord instance (for the `compare (C x) (C y)` case).
instance (Ord a) => Ord (Elem2 a b) where
    compare (Val _) (Val _)   = EQ
    compare (Val _) x         = LT
    compare x        (Val _)  = GT
    compare (C x) (C y)       = compare x y

isVal (Val _) = True
isVal _       = False

-- | 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 = TNode (C x) TEnd (singleton xs b) TEnd
singleton []     b = TNode (Val b) TEnd TEnd TEnd

-- | 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
-- General case
insert xss@(x:xs) b (TNode ele l e h) =
    case compare (C x) ele of
        LT -> TNode ele (insert xss b l) e h
        EQ -> TNode ele l (insert xs b e) h
        GT -> TNode ele l e (insert xss b h)
-- Insert new elements quickly
insert xss@(x:xs) b TEnd =
    singleton xss b
-- end of word in non empty tree
insert [] b (TNode ele l e h) = 
    case compare (Val b) ele of
        EQ -> TNode (Val b) l e h
        LT  -> TNode ele (insert [] b l) e h
-- end of word in empty tree
insert [] b TEnd =
    TNode (Val b) TEnd TEnd TEnd


-- | Returns true if the `[a]` is a key in the TernaryMap.
member :: Ord a => [a] -> TernaryMap a b -> Bool
member          _ TEnd              = False
member         [] (TNode ele l e h) = isVal ele || member [] l
member xss@(x:xs) (TNode ele l e h) = 
    case compare (C 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 _ TEnd = Nothing
lookup [] (TNode (Val b) _ _ _) = Just b
lookup [] (TNode ele l _ _)     = lookup [] l
lookup xss@(x:xs) (TNode ele l e h) =
    case compare (C 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 TEnd = 0
treeSize (TNode (Val _) l e h) = treeSize l + treeSize e + treeSize h
treeSize (TNode _ l e h) = 1 + treeSize l + treeSize e + treeSize h

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

-- | 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) TEnd

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

-- | Makes a list of all the values in the map.
elems :: TernaryMap a b -> [b]
elems (TEnd) = []
elems (TNode (Val v) l _ h) = elems l ++ (v : elems h)
elems (TNode _ l e h)       = elems l ++ (elems e ++ elems h)

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

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

instance Functor (Elem2 a) where
    fmap _ (C x) = C x
    fmap f (Val b) = Val . f $ b

instance Functor (TernaryMap a) where
    fmap f (TNode ele l e h)     = TNode (fmap f ele) (fmap f l) (fmap f e) (fmap f h)
    fmap _ TEnd = TEnd

instance (Binary a, Binary b) => Binary (Elem2 a b) where
    put (C a) = putWord8 0 >> put a
    put (Val b) = putWord8 1 >> put b
    get = do
        n <- getWord8
        case n of
            0 -> liftM C get
            1 -> liftM Val get

-- | A rather long Binary instance, that uses binary numbers to indicate
-- where TEnds are efficiently.
instance (Binary a, Binary b) => Binary (TernaryMap a b) where
    put (TNode ch TEnd TEnd TEnd) = do
        putWord8 0
        put ch
    put (TNode ch TEnd TEnd h) = do
        putWord8 1
        put ch
        put h
    put (TNode ch TEnd e TEnd) = do
        putWord8 2
        put ch
        put e
    put (TNode ch TEnd e h) = do
        putWord8 3
        put ch
        put e
        put h
    put (TNode ch l TEnd TEnd) = do
        putWord8 4
        put ch
        put l
    put (TNode ch l TEnd h) = do
        putWord8 5
        put ch
        put l
        put h
    put (TNode ch l e TEnd) = do
        putWord8 6
        put ch
        put l
        put e
    -- General case
    put (TNode ch l e h) = do
        putWord8 7
        put ch
        put l
        put e
        put h
    put TEnd = putWord8 8
    get = do
        tag <- getWord8
        case tag of
            8 -> return TEnd
            _ -> do
                ch <- get
                case tag of
                    0 -> return (TNode ch TEnd TEnd TEnd)
                    1 -> do
                        h <- get
                        return (TNode ch TEnd TEnd h)
                    2 -> do
                        e <- get
                        return (TNode ch TEnd e TEnd)
                    3 -> do
                        e <- get
                        h <- get
                        return (TNode ch TEnd e h)
                    4 -> do
                        l <- get
                        return (TNode ch l TEnd TEnd)
                    5 -> do
                        l <- get
                        h <- get
                        return (TNode ch l TEnd h)
                    6 -> do
                        l <- get
                        e <- get
                        return (TNode ch l e TEnd)
                    7 -> do
                        l <- get
                        e <- get
                        h <- get
                        return (TNode ch l e h)