module Data.Binary.Defer.Trie(
    Trie, newTrie, lookupTrie, unionsWith
    ) where

import Data.Binary.Defer
import Data.List
import Data.Maybe
import General.Code


-- Keep simple first, can move to an Array later if necessary
-- the second list is always ordered
data Trie a = Trie (Maybe a) [(Char,Trie a)]
              deriving Eq


newTrie :: [(String,a)] -> Trie a
newTrie = newTrieOrdered . sortBy (compare `on` fst)


newTrieOrdered :: [(String,a)] -> Trie a
newTrieOrdered xs = Trie
        (fmap snd $ listToMaybe as)
        (map f $ groupBy ((==) `on` (head . fst)) bs)
    where
        (as,bs) = span (null . fst) xs
        f xs = (head $ fst $ head xs, newTrieOrdered [(a,b) | (_:a,b) <- xs])


lookupTrie :: String -> Trie a -> Maybe a
lookupTrie [] (Trie a b) = a
lookupTrie (x:xs) (Trie a b) = lookup x b >>= lookupTrie xs


unionsWith :: (a -> a -> a) -> [Trie a] -> Trie a
unionsWith f [] = newTrie []
unionsWith f [x] = x
unionsWith f (x:xs) = foldl (unionWith f) x xs


unionWith :: (a -> a -> a) -> Trie a -> Trie a -> Trie a
unionWith merge (Trie x xs) (Trie y ys) = Trie xy $ f xs ys
    where
        xy = case (x,y) of
                  (Just x, Nothing) -> Just x
                  (Nothing, Just y) -> Just y
                  (Just x,  Just y) -> Just $ merge x y
                  (Nothing,Nothing) -> Nothing

        f ((x1,x2):xs) ((y1,y2):ys) = case compare x1 y1 of
            EQ -> (x1,unionWith merge x2 y2) : f xs ys
            LT -> (x1,x2) : f xs ((y1,y2):ys)
            GT -> (y1,y2) : f ((x1,x2):xs) ys
        f xs [] = xs
        f [] ys = ys


instance BinaryDefer a => BinaryDefer (Trie a) where
    put (Trie a b) = putDefer $ put2 a b
    get = getDefer $ get2 Trie

instance Show a => Show (Trie a) where
    show = unlines . f ""
        where
            f str (Trie res xs) =
                [show str ++ " = " ++ show r | Just r <- [res]] ++
                concat [f (str ++ [c]) t | (c,t) <- xs]