{-# LANGUAGE BangPatterns, UnboxedTuples, TupleSections, TypeFamilies, PatternGuards, UnboxedTuples #-}

module Data.TrieMap.RadixTrie () where

import Data.TrieMap.TrieKey
import Data.TrieMap.Sized
import Data.TrieMap.Applicative

import Control.Applicative
import Control.Monad

import Data.Maybe
import Data.Foldable
import Data.Traversable

import Prelude hiding (lookup, foldr, foldl)

data Edge k m a = Edge {-# UNPACK #-} !Int [k] (Maybe a) (m (Edge k m a))
type Edge' k a = Edge k (TrieMap k) a
type MEdge' k a = Maybe (Edge' k a)

edgeSize :: Edge k m a -> Int
edgeSize (Edge sz _ _ _) = sz

instance TrieKey k =>  TrieKey [k] where
	newtype TrieMap [k] a = Radix (MEdge' k a)
	emptyM = Radix Nothing
	singletonM s ks a = Radix (Just (Edge (s a) ks (Just a) emptyM))
	nullM (Radix m) = isNothing m
	sizeM _ (Radix m) = maybe 0 edgeSize m
	lookupM ks (Radix m) = m >>= lookup ks
	alterM s f ks (Radix m) = Radix (alter s f ks m)
	alterLookupM s f ks (Radix m) = onUnboxed Radix (alterLookupE s f ks) m
	traverseWithKeyM s f (Radix m) = Radix <$> traverse (traverseE s f) m
	extractM s f (Radix m) = maybe empty (fmap Radix <.> extractE s f) m
	foldWithKeyM f (Radix m) z = foldr (foldE f) z m
	foldlWithKeyM f (Radix m) z = foldl (foldlE f) z m
	mapMaybeM s f (Radix m) = Radix (m >>= mapMaybeE s f)
	mapEitherM _ _ _ (Radix Nothing) = (# emptyM, emptyM #)
	mapEitherM s1 s2 f (Radix (Just m)) = both Radix Radix (mapEitherE s1 s2 f) m
	unionM s f (Radix m1) (Radix m2) = Radix (unionMaybe (unionE s f) m1 m2)
	isectM s f (Radix m1) (Radix m2) = Radix (isectMaybe (isectE s f) m1 m2)
	diffM s f (Radix m1) (Radix m2) = Radix (diffMaybe (diffE s f) m1 m2)
-- 	lookupIxM s ks (Radix m) = maybe (empty, empty, empty) (lookupIxE s 0 ks) m
	isSubmapM (<=) (Radix m1) (Radix m2) = subMaybe (isSubmapE (<=)) m1 m2
	splitLookupM _ _ _ (Radix Nothing) = (# emptyM, Nothing, emptyM #)
	splitLookupM s f ks (Radix (Just e)) = sides Radix (splitLookupE s f ks) e
-- 	assocAtM s i (Radix m) = maybe (empty, empty, empty) (assocAtE s 0 i) m
  
cat :: [k] -> Edge' k a -> Edge' k a
ks `cat` Edge sz ls v ts = Edge sz (ks ++ ls) v ts

cons :: k -> Edge' k a -> Edge' k a
k `cons` Edge sz ks v ts = Edge sz (k:ks) v ts

edge :: TrieKey k =>  Sized a -> [k] -> Maybe a -> TrieMap k (Edge' k a) -> Edge' k a
edge s ks v ts = Edge (maybe 0 s v + sizeM edgeSize ts) ks v ts

singleMaybe :: TrieKey k => Sized a -> [k] -> Maybe a -> MEdge' k a
singleMaybe s ks v = do	v <- v
			return (edge s ks (Just v) emptyM)

compact :: TrieKey k => Edge' k a -> MEdge' k a
compact e@(Edge _ ks Nothing ts) = case assocsM ts of
	[]	-> Nothing
	[(l, e')] -> compact (ks `cat` (l `cons` e'))
	_	-> Just e
compact e = Just e

lookup :: (Eq k, TrieKey k) => [k] -> Edge' k a -> Maybe a
lookup ks (Edge _ ls v ts) = match ks ls where
	match (k:ks) (l:ls)
		| k == l = match ks ls
	match (k:ks) [] = lookupM k ts >>= lookup ks
	match [] [] = v
	match _ _ = Nothing

alter :: TrieKey k => Sized a -> (Maybe a -> Maybe a) -> [k] -> MEdge' k a -> MEdge' k a
alter s f ks0 Nothing = singleMaybe s ks0 (f Nothing)
alter s f ks0 (Just e@(Edge sz ls0 v ts)) = match 0 ks0 ls0 where
	match !i (k:ks) (l:ls) = case compare k l of
	      LT | Just v' <- f Nothing	
		      -> Just $ let sv = s v' in Edge (sv + sz) (take i ls0) Nothing (fromDistAscListM edgeSize
					[(k, Edge sv ks (Just v') emptyM), (l, Edge sz ls v ts)])
	      EQ	-> match (i+1) ks ls
	      GT | Just v' <- f Nothing
		      -> Just $ let sv = s v' in Edge (sv + sz) (take i ls0) Nothing (fromDistAscListM edgeSize
					[(l, Edge sz ls v ts), (k, Edge sv ks (Just v') emptyM)])
	      _	-> Just e
	match _ (k:ks) [] = compact $ edge s ls0 v (alterM edgeSize g k ts) where
		g = alter s f ks
	match _ [] (l:ls)
		| Just v' <- f Nothing
			= Just (Edge (s v' + sz) ks0 (Just v') (singletonM edgeSize l (Edge sz ls v ts)))
	match _ [] []
		= compact (edge s ls0 (f v) ts)
	match _ _ _ = Just e

alterLookupE :: TrieKey k => Sized a -> (Maybe a -> (# z, Maybe a #)) -> [k] -> MEdge' k a -> (# z, MEdge' k a #)
alterLookupE s f ks Nothing = onUnboxed (singleMaybe s ks) f Nothing
alterLookupE s f ks0 (Just e@(Edge sz ls0 v0 ts0)) = match 0 ks0 ls0 where
      match !i (k:ks) (l:ls) = case compare k l of
	      LT	-> onUnboxed (Just . maybe e (\ v' -> let sv = s v' in Edge (sz + sv) (take i ls0) Nothing $
				      fromDistAscListM edgeSize [(k, Edge sv ks (Just v') emptyM), (l, Edge sz ls v0 ts0)]))
			      f Nothing
	      GT	-> onUnboxed (Just . maybe e (\ v' -> let sv = s v' in Edge (sz + sv) (take i ls0) Nothing $
				      fromDistAscListM edgeSize [(l, Edge sz ls v0 ts0), (k, Edge sv ks (Just v') emptyM)]))
			      f Nothing
	      EQ	-> match (i+1) ks ls
      match _ (k:ks) [] = onUnboxed (compact . edge s ls0 v0) (alterLookupM edgeSize g k) ts0 where
	      g = alterLookupE s f ks
      match _ [] (l:ls) = onUnboxed (Just . maybe e (\ v' -> let sv = s v' in 
					Edge (sv + sz) ks0 (Just v') (singletonM edgeSize l (Edge sz ls v0 ts0))))
			      f Nothing
      match _ [] [] = onUnboxed (\ v' -> compact $ edge s ls0 v' ts0) f v0

traverseE :: (Applicative f, TrieKey k) => Sized b -> ([k] -> a -> f b) -> Edge' k a -> f (Edge' k b)
traverseE s f (Edge _ ks v ts)
	= edge s ks <$> traverse (f ks) v <*> traverseWithKeyM edgeSize g ts 
	where	g l = traverseE s (\ ls -> f (ks ++ l:ls))

extractE :: (Alternative f, TrieKey k) => Sized a -> ([k] -> a -> f (x, Maybe a)) -> Edge' k a -> f (x, MEdge' k a)
extractE s f (Edge _ ks v ts) = case v of
	Nothing	-> rest
	Just v	-> fmap (\ v' -> compact (edge s ks v' ts)) <$> f ks v <|> rest
	where	rest = fmap (compact . edge s ks v) <$> extractM edgeSize g ts
	     	g l = extractE s (\ ls -> f (ks ++ l:ls))

foldE :: TrieKey k => ([k] -> a -> b -> b) -> Edge' k a -> b -> b
foldE f (Edge _ ks v ts) z = foldr (f ks) (foldWithKeyM g ts z) v where
	g l = foldE (\ ls -> f (ks ++ l:ls))

foldlE :: TrieKey k => ([k] -> b -> a -> b) -> b -> Edge' k a -> b 
foldlE f z (Edge _ ks v ts) = foldlWithKeyM g ts (foldl (f ks) z v) where
	g l = foldlE (\ ls -> f (ks ++ l:ls))

mapMaybeE :: TrieKey k => Sized b -> ([k] -> a -> Maybe b) -> Edge' k a -> MEdge' k b
mapMaybeE s f (Edge _ ks v ts) = compact (edge s ks (v >>= f ks)
	(mapMaybeM edgeSize (\ l -> mapMaybeE s (\ ls -> f (ks ++ l:ls))) ts))

mapEitherE :: TrieKey k => Sized b -> Sized c -> ([k] -> a -> (# Maybe b, Maybe c #)) -> Edge' k a ->
	(# MEdge' k b, MEdge' k c #)
mapEitherE s1 s2 f (Edge _ ks v ts) = case mapEitherM edgeSize edgeSize (\ l -> mapEitherE s1 s2 (\ ls -> f (ks ++ l:ls))) ts of
  (# tsL, tsR #) -> case v of
       Nothing	-> (# compact (edge s1 ks Nothing tsL), compact (edge s2 ks Nothing tsR) #)
       Just v	-> case f ks v of
		      (# vL, vR #) -> (# compact (edge s1 ks vL tsL), compact (edge s2 ks vR tsR) #)

unionE :: TrieKey k => Sized a -> ([k] -> a -> a -> Maybe a) -> Edge' k a -> Edge' k a -> MEdge' k a
unionE s f (Edge szK ks0 vK tsK) (Edge szL ls0 vL tsL) = match 0 ks0 ls0 where
	match !i (k:ks) (l:ls) = case compare k l of
	      EQ -> match (i+1) ks ls
	      LT -> Just $ Edge (szK + szL) (take i ks0) Nothing (fromDistAscListM edgeSize 
		      [(k, Edge szK ks vK tsK), (l, Edge szL ls vL tsL)])
	      GT -> Just $ Edge (szK + szL) (take i ks0) Nothing (fromDistAscListM edgeSize 
		      [(l, Edge szL ls vL tsL), (k, Edge szK ks vK tsK)])
	match _ [] (l:ls) = compact (edge s ks0 vK (alterM edgeSize g l tsK)) where
		g (Just eK') = unionE s (\ ls' -> f (ks0 ++ l:ls')) eK' (Edge szL ls vL tsL)
		g Nothing = Just (Edge szL ls vL tsL)
	match _ (k:ks) [] = compact (edge s ls0 vL (alterM edgeSize g k tsL)) where
		g Nothing = Just (Edge szK ks vK tsK)
		g (Just eL') = unionE s (\ ks' -> f (ls0 ++ k:ks')) (Edge szK ks vK tsK) eL'
	match _ [] [] = compact (edge s ls0 (unionMaybe (f ls0) vK vL) (unionM edgeSize g tsK tsL)) where
		g x = unionE s (\ xs -> f (ls0 ++ x:xs))

isectE :: TrieKey k => Sized c -> ([k] -> a -> b -> Maybe c) -> Edge' k a -> Edge' k b -> MEdge' k c
isectE s f (Edge szK ks0 vK tsK) (Edge szL ls0 vL tsL) = match ks0 ls0 where
	match (k:ks) (l:ls)
		| k == l	= match ks ls
	match (k:ks) [] = do	eL' <- lookupM k tsL
			   	cat ls0 <$> cons k <$> isectE s (\ ks' -> f (ls0 ++ k:ks')) (Edge szK ks vK tsK) eL'
	match [] (l:ls) = do	eK' <- lookupM l tsK
				cat ks0 <$> cons l <$> isectE s (\ ls' -> f (ks0 ++ l:ls')) eK' (Edge szL ls vL tsL)
	match [] [] = compact (edge s ks0 (isectMaybe (f ks0) vK vL) (isectM edgeSize g tsK tsL)) where
		g x = isectE s (\ xs -> f (ks0 ++ x:xs))
	match _ _ = Nothing

diffE :: TrieKey k => Sized a -> ([k] -> a -> b -> Maybe a) -> Edge' k a -> Edge' k b -> MEdge' k a
diffE s f eK@(Edge szK ks0 vK tsK) (Edge szL ls0 vL tsL) = match ks0 ls0 where
	match (k:ks) (l:ls)
		| k == l	= match ks ls
	match (k:ks) []
		| Just eL' <- lookupM k tsL
			= cat ls0 . cons k <$> diffE s (\ ks' -> f (ls0 ++ k:ks')) (Edge szK ks vK tsK) eL'
	match [] (l:ls)
		= compact (edge s ks0 vK (alterM edgeSize (>>= g) l tsK))
		where	g eK' = diffE s (\ ls' -> f (ks0 ++ l:ls')) eK' (Edge szL ls vL tsL)
	match [] [] = compact (edge s ks0 (diffMaybe (f ks0) vK vL) (diffM edgeSize g tsK tsL)) where
		g x = diffE s (\ xs -> f (ks0 ++ x:xs))
	match _ _ = Just eK


isSubmapE :: TrieKey k => LEq a b -> LEq (Edge' k a) (Edge' k b)
isSubmapE (<=) (Edge szK ks vK tsK) (Edge _ ls vL tsL) = match ks ls where
	match (k:ks) (l:ls)
		| k == l	= match ks ls
	match (k:ks) []
		| Just eL' <- lookupM k tsL
			= isSubmapE (<=) (Edge szK ks vK tsK) eL'
	match [] [] = subMaybe (<=) vK vL && isSubmapM (isSubmapE (<=)) tsK tsL
	match _ _ = False

splitLookupE :: TrieKey k => Sized a -> (a -> (# Maybe a, Maybe x, Maybe a #)) -> [k] -> Edge' k a ->
	(# MEdge' k a, Maybe x, MEdge' k a #)
splitLookupE s f ks e@(Edge _ ls v ts) = match ks ls where
	match (k:ks) (l:ls) = case compare k l of
		LT	-> (# Nothing, Nothing, Just e #)
		GT	-> (# Just e, Nothing, Nothing #)
		EQ	-> match ks ls
	match (k:ks) [] = case splitLookupM edgeSize g k ts of
	    (# tsL, x, tsR #) -> (# compact (edge s ls v tsL), x, compact (edge s ls Nothing tsR) #)
	  where	g = splitLookupE s f ks
	match [] (_:_) = (# Nothing, Nothing, Just e #)
	match [] [] = case v of
	    Nothing	-> (# Nothing, Nothing, compact (edge s ls Nothing ts) #)
	    Just v	-> case f v of
		(# vL, x, vR #)	-> (# singleMaybe s ls vL, x, compact (edge s ls vR ts) #)