{-# OPTIONS -Wno-incomplete-uni-patterns #-}
{-# LANGUAGE TemplateHaskell #-}

module SimpleTrie where

import Data.Monoid
import Control.Lens
import qualified Data.Map.Strict as M

data Trie a b = Trie
    { _trieMVal :: Maybe b
    , _trieChildren :: M.Map a (Trie a b)
    }

makeLenses ''Trie

emptyTrie :: Trie a b
emptyTrie = Trie Nothing M.empty

singletonTrie :: Ord a => [a] -> b -> Trie a b
singletonTrie (k : ks) x = Trie Nothing $ M.singleton k $ singletonTrie ks x
singletonTrie [] x = Trie (Just x) M.empty

insertInTrie :: Ord a => [a] -> b -> Trie a b -> Trie a b
insertInTrie [] x = trieMVal ?~ x
insertInTrie (k : ks) x = trieChildren %~ M.insertWith f k (singletonTrie ks x)
  where
    f _ = insertInTrie ks x

lookupTrie :: Ord a => [a] -> Trie a b -> Maybe b
lookupTrie (k : ks) t = _trieChildren t M.!? k >>= lookupTrie ks
lookupTrie [] t = _trieMVal t

----nextTrie :: Ord a => [a] -> Trie a b -> Maybe (b, Trie a b)
--nextTrie :: [a] -> Trie a b -> Maybe (b, Trie a b)
--nextTrie (k:ks) t = undefined
--nextTrie [] t = Nothing

splitLookupTrie :: Ord a => [a] -> Trie a b -> (Trie a b,Maybe b,Trie a b)
splitLookupTrie (k:ks) (Trie x xs) = case m of
    Nothing -> (Trie x l,Nothing,Trie x r)
    Just t' -> let (l',m',r') = splitLookupTrie ks t'
        in (Trie x (M.insert k l' l),m', Trie x (M.insert k r' r))
  where
    (l,m,r) = M.splitLookup k xs
splitLookupTrie [] (Trie x xs) = (Trie Nothing mempty,x, Trie Nothing xs)
-- the above would be better if it removed empty trie strings; the left trie
-- will always be empty and the right trie might be empty.
-- I cannot think of a smart way to do this

firstTrie :: Ord a => Trie a b -> Maybe b
firstTrie t = case t ^. trieMVal of
    Just x -> Just x
    Nothing -> getFirst $ foldMap (First . firstTrie) (t ^. trieChildren)

-- # OPTIONS -Wno-incomplete-uni-patterns #-}
multiLookupTrie :: Ord a => [a] -> Trie a b -> [([a], b)]
multiLookupTrie xs t@(Trie my ch)
    | null xs || null ch = maybe [] (\y -> [([], y)]) my
    | otherwise =
        let (z : zs) = xs
         in case M.lookup z ch of
                Nothing -> multiLookupTrie zs t
                Just t' -> map (_1 %~ (z :)) (multiLookupTrie zs t') ++ multiLookupTrie zs t

multiLookupTrieI :: Ord a => [(a, c)] -> Trie a b -> [([c], b)]
multiLookupTrieI xs t@(Trie my ch)
    | null xs || null ch = maybe [] (\y -> [([], y)]) my
    | otherwise =
        let ((z1, z2) : zs) = xs
         in case M.lookup z1 ch of
                Nothing -> multiLookupTrieI zs t
                Just t' -> map (_1 %~ (z2 :)) (multiLookupTrieI zs t') ++ multiLookupTrieI zs t