module Trie (insert, insertWith, singleton, empty, lookup --)
             ,Trie (..) )
    where



import Prelude hiding (lookup)


 --- MODULE TO USE INTERNALLY TO STORE BRANCHES (pick one): ---
  --- FASTEST:
import qualified MyMap  as M       -- unbalanced tree
  --- SLOWER:
--import qualified DataMap  as M  --Data.Map with strict value field
--import qualified Data.Map  as M  
  --- SLOWEST:
--import qualified MyListMap  as M   -- simple association lists


data Trie a v = Node    { branches :: !(M.Map a (Trie a v)) }
              | ValNode { branches :: !(M.Map a (Trie a v)),
                          val :: v }         --mid-word, and last letter of word
              | ValBucket { bucket :: ![a],  --remainder of a word
                           val :: v  }
              | Val { val :: v }
              | Empty
               deriving (Show)


insertWith :: (Ord a)=> (v->v->v) -> [a] -> v -> Trie a v -> Trie a v
insertWith _ as v Empty = singleton as v

insertWith _ [] v (Node m) = ValNode m v
 -- insert value mid-word: e.g "was" --> "washing"
insertWith f [] v (ValBucket as v') = ValNode (intoMap as v') v 
insertWith f [] v vn = vn{ val = f v (val vn) }  -- ValNode, Val

insertWith f n v (Val v') = ValNode (intoMap n v) v'
 -- make ValBucket into a Node and insertWith into that:
insertWith f n v (ValBucket o v') = insertWith f n v $Node (intoMap o v') 
insertWith f (a:as) v n =   -- Node, ValNode
    n{ branches = 
        --- ONLY WITH Data.Map, DataMap, MyMap:
        --M.insertWith' (\_->insertWith f as v) a (singleton as v)$branches n }
        M.insertWith (\_->insertWith f as v) a (singleton as v)$branches n }


-- convert a bucket/value pair into a singleton branch Map :
intoMap (a:as) v = M.singleton a$ singleton as v


-- merge the key-list/value pair into the Trie
insert :: (Ord a)=> [a] -> v -> Trie a v -> Trie a v
insert = insertWith const

singleton :: (Ord a)=> [a] -> v -> Trie a v
singleton [] = Val
singleton l  = ValBucket l

empty :: (Ord a)=> Trie a v
empty = Empty



lookup :: (Ord a, Monad m)=> [a] -> Trie a v -> m v
lookup _  Empty    = fail "empty Trie"
lookup [] (ValNode _ v) = return v
lookup [] (Val v)       = return v
lookup []  _      = fail "not present"
lookup as (Val _) = fail "not present"
lookup as (ValBucket as' v)| as == as' = return v
                           | otherwise = fail "not present"
lookup (a:as) n = M.lookup a (branches n) >>= lookup as  -- ValNode and Node