{-# LANGUAGE IncoherentInstances, MultiParamTypeClasses, UndecidableInstances, FlexibleContexts, StandaloneDeriving, PatternGuards #-}

module TrieMap.RadixTrie (RadixTrie) where

import Control.Applicative hiding (Alternative(..))
import Control.Monad
import Data.Foldable
import Data.Traversable
import Data.Monoid
import Data.Maybe
import Data.Ord
import Data.Sequence (Seq, (|>))
import qualified Data.Sequence as Seq

import TrieMap.MapTypes
import TrieMap.TrieAlgebraic
import TrieMap.Applicative

import Prelude hiding (null, foldr, all)

instance (Eq k, Eq v, TrieKey k m) => Eq (Edge k m v) where
	Edge n1 ks1 v1 ts1 == Edge n2 ks2 v2 ts2 = n1 == n2 && ks1 == ks2 && v1 == v2 && assocsAlg ts1 == assocsAlg ts2

instance (Ord k, Ord v, TrieKey k m) => Ord (Edge k m v) where
	Edge _ ks1 v1 ts1 `compare` Edge _ ks2 v2 ts2 = 
		compare ks1 ks2 `mappend` compare v1 v2 `mappend` comparing assocsAlg ts1 ts2

deriving instance (Eq k, Eq v, TrieKey k m) => Eq (RadixTrie k m v)
deriving instance (Ord k, Ord v, TrieKey k m) => Ord (RadixTrie k m v)
deriving instance (Show k, Show v, Functor m, Show (m String)) => Show (RadixTrie k m v)

instance (Show k, Show v, Functor m, Show (m String)) => Show (Edge k m v) where
	show (Edge _ k v ts) = "Edge " ++ show k ++ " " ++ show v ++ " " ++ show (fmap show ts)

instance Sized (Edge k m v) where
	getSize (Edge n _ _ _) = n

instance (Ord k, TrieKey k m) => TrieKey [k] (RadixTrie k m) where
	emptyAlg = Radix Nothing
	nullAlg = isNothing . unRad
	sizeAlg (Radix e) = maybe 0 getSize e
	getSingleAlg (Radix e) = e >>= getSingleEdge
	guardNullAlg (Radix e) = do	e <- guardNullEdge =<< e
					return (Radix (Just e))
	lookupAlg ks = unRad >=> lookupEdge (==) ks
-- 	sizeAlg (Radix e) = maybe 0 sizeEdge e
	alterLookupAlg f k = fmap Radix .
		maybe (fmap (maybeSingleEdge k) $ f Nothing)
			(alterLookupEdge (==) f k) . unRad
	foldWithKeyAlg f z = foldr (flip (foldWithKeyEdge f)) z . unRad
	mapMaybeAlg f (Radix e) = Radix (e >>= mapMaybeEdge f)
	mapEitherAlg f (Radix Nothing) = (emptyAlg, emptyAlg)
	mapEitherAlg f (Radix (Just e)) = (Radix e1, Radix e2)
		where	(e1, e2) = mapEitherEdge f e
-- 	mapMaybeAlg f (Radix e) = (Radix . join) <$> traverse (mapAppMaybeEdge f) e
	mapAppAlg f = fmap Radix . traverse (mapAppEdge f) . unRad
	unionMaybeAlg f (Radix e1) (Radix e2) = Radix (unionMaybe (unionMaybeEdge f) e1 e2)
	intersectAlg f (Radix e1) (Radix e2) = Radix (intersectMaybe (intersectEdge f) e1 e2)
	differenceAlg f (Radix e1) (Radix e2) = Radix (differenceMaybe (differenceEdge f) e1 e2)

	getMinAlg (Radix e) = fmap (fmap Radix . getMinEdge) e
	getMaxAlg (Radix e) = fmap (fmap Radix . getMaxEdge) e
-- 	updateMinAlg f (Radix e) = Radix $ e >>= updateMinEdge f
-- 	updateMaxAlg f (Radix e) = Radix $ e >>= updateMaxEdge f

	fromListAlg f xs = Radix (edgeFromList f xs)
	fromAscListAlg f xs = Radix (edgeFromAscList f xs)
	fromDistAscListAlg = fromAscListAlg (\ _ v _ -> v)

	isSubmapAlg (<=) (Radix e1) (Radix e2) = isSubmapAlg subEdge e1 e2 -- hehe, using the Maybe instance here!
		where	subEdge = isSubmapEdge (==) (<=) lookupAlg $! isSubmapAlg subEdge

	valid (Radix e) = maybe True validEdge e

	splitLookupAlg _ _ (Radix Nothing) = (Radix Nothing, Nothing, Radix Nothing)
	splitLookupAlg f k (Radix (Just e)) = case splitEdge f k e of
		(eL, ans, eR)	-> (Radix eL, ans, Radix eR)

-- sizeEdge :: Edge k m v -> Int
-- sizeEdge (Edge n _ _ _) = n

{-# INLINE edge #-}
edge :: (Sized v, TrieKey k m) => [k] -> Maybe v -> m (Edge k m v) -> Edge k m v
edge ks v ts = Edge (getSize v + getSize ts) ks v ts

lookupEdge :: TrieKey k m => (k -> k -> Bool) -> [k] -> Edge k m v -> Maybe v
lookupEdge (==) ks (Edge _ ls v ts) = procEdge ks ls where
	procEdge (k:ks) (l:ls)
		| k == l  = procEdge ks ls
	procEdge (k:ks) [] = lookupAlg k ts >>= lookupEdge (==) ks
	procEdge [] [] = v
	procEdge _ _ = Nothing

edgeFromList :: (Eq k, TrieKey k m, Sized v) => ([k] -> v -> v -> v) -> [([k], v)] -> MEdge k m v
edgeFromList f xs = guardNullEdge $ edge [] v0 $ mapMaybeAlg (\ k (Elem xs)-> edgeFromList (f . (k:)) xs) $ 
	fromListAlg (\ _ (Elem xs) (Elem ys) -> Elem (ys ++ xs)) ys
	where	part ([], v) (v0, ys) = (Just $ maybe v (flip (f []) v) v0, ys)
		part (k:ks, v) (v0, ys) = (v0, (k, Elem [(ks, v)]):ys)
		(v0, ys) = foldr part (Nothing, []) xs

edgeFromAscList :: (Eq k, TrieKey k m, Sized v) => ([k] -> v -> v -> v) -> [([k], v)] -> MEdge k m v
edgeFromAscList _ [] = Nothing
edgeFromAscList f (x:xs) = Just $ edgeFromAscList' f x xs

edgeFromAscList' :: (Eq k, TrieKey k m, Sized v) => ([k] -> v -> v -> v) -> ([k], v) -> [([k], v)] -> Edge k m v
edgeFromAscList' f (ks, v) [] = Edge (getSize v) ks (Just v) emptyAlg
edgeFromAscList' f x xs = case groupHead f (x:xs) of
	(Nothing, [(k, ~(Edge n ks v ts))])
			-> Edge n (k:ks) v ts
	(ans, xs')	-> edge [] ans (fromDistAscListAlg xs')

groupHead :: (Eq k, TrieKey k m, Sized v) => ([k] -> v -> v -> v) -> [([k], v)] -> (Maybe v, [(k, Edge k m v)])
groupHead f (([], v):xs) = case groupHead f xs of
	(v', ans)	-> (Just $ maybe v (f [] v) v', ans)
groupHead f ((k:ks, v):xs) = (Nothing, groupHead' k (ks, v) Seq.empty xs) where
	groupHead' k0 x xs ((k:ks, v):ys)
		| k == k0	= groupHead' k0 x (xs |> (ks, v)) ys
		| otherwise	= (k0, edgeFromAscList' (f . (k0:)) x (toList xs)):groupHead' k (ks, v) Seq.empty ys
	groupHead' k0 x xs [] = [(k0, edgeFromAscList' (f . (k0:)) x (toList xs))]
	groupHead' _ _ _ _ = error "Violation of ascending invariant!"
groupHead _ [] = (Nothing, [])
 {-guardNullEdge $ Edge [] v0 $ mapMaybeAlg (\ k -> edgeFromAscList (f . (k:))) $ fromAscListAlg (const (flip (++))) ys
	where	part ([], v) (v0, ys) = (Just $ maybe v (flip (f []) v) v0, ys)
		part (k:ks, v) (v0, ys) = (v0, (k, [(ks, v)]):ys)
		(v0, ys) = foldr part (Nothing, []) xs-}

maybeSingleEdge :: Sized v => TrieKey k m => [k] -> Maybe v -> MEdge k m v
maybeSingleEdge ks = fmap (\ v -> Edge (getSize v) ks (Just v) emptyAlg)

getSingleEdge :: (TrieKey k m) => Edge k m v -> Maybe ([k], v)
getSingleEdge (Edge _ ks (Just v) ts)
	| nullAlg ts	= Just (ks, v)
getSingleEdge (Edge _ ks Nothing ts) = do
	(x, e') <- getSingleAlg ts
	(xs, v) <- getSingleEdge e'
	return (ks ++ x:xs, v) 
getSingleEdge _ = Nothing

{-# INLINE guardNullEdge #-}
guardNullEdge :: TrieKey k m => Edge k m v -> MEdge k m v
guardNullEdge (Edge n ks Nothing ts)
	| nullAlg ts	= Nothing
	| Just (x, Edge n' xs v ts') <- getSingleAlg ts
		= Just (Edge n' (ks ++ x:xs) v ts')
guardNullEdge e = Just e

alterLookupEdge :: (TrieKey k m, Sized v) => (k -> k -> Bool) ->
	(Maybe v -> (a, Maybe v)) -> [k] -> Edge k m v -> (a, MEdge k m v)
alterLookupEdge (==) f ks0 e@(Edge n0 ls0 v ts) = procEdge 0 ks0 ls0 where
	procEdge i _ _ | i `seq` False = undefined
	procEdge i (k:ks) (l:ls)
		| k == l	= procEdge (i+1) ks ls
		| otherwise	= fmap (Just . g) (f Nothing)
		where	g Nothing = e
			g (Just v') = let nV = getSize v' in Edge (n0 + nV) (take i ks0) Nothing $
						fromListAlg' [(k, Edge nV ks (Just v') emptyAlg), (l, Edge n0 ls v ts)]
	procEdge i (k:ks) [] = proc (alterLookupAlg g k ts) where
		g Nothing = maybeSingleEdge ks <$> f Nothing
		g (Just e') = alterLookupEdge (==) f ks e'
		proc = fmap (guardNullEdge . edge ls0 v)
	procEdge i [] (l:ls) = fmap (Just . g) $ f Nothing
		where	g Nothing = e
			g (Just v') = Edge (getSize v' + n0) ks0 (Just v') $ insertAlg l (Edge n0 ls v ts) emptyAlg
	procEdge i [] [] = (ans, guardNullEdge (Edge (getSize fv - getSize v + n0) ks0 fv ts))
		where	(ans, fv) = f v

foldWithKeyEdge :: TrieKey k m => ([k] -> v -> x -> x) -> x -> Edge k m v -> x
foldWithKeyEdge f z (Edge _ ks v ts) =
	foldr (f ks) (foldWithKeyAlg (\ x -> flip (foldWithKeyEdge (\ xs -> f (ks ++ x:xs)))) z ts) v

mapMaybeEdge :: (TrieKey k m, Sized w) => ([k] -> v -> Maybe w) -> Edge k m v -> MEdge k m w
mapMaybeEdge f (Edge _ ks v ts) = guardNullEdge $
	edge ks (join $ traverse (f ks) v) (mapMaybeAlg (\ x -> mapMaybeEdge (\ xs -> f (ks ++ x:xs))) ts)

mapEitherEdge :: (TrieKey k m, Sized b, Sized c) => ([k] -> a -> Either b c) -> Edge k m a -> (MEdge k m b, MEdge k m c)
mapEitherEdge f (Edge _ ks v ts) =
	(guardNullEdge $ edge ks vL tsL, guardNullEdge $ edge ks vR tsR) 
	where	(vL, vR) = case fmap (f ks) v of
			Nothing	-> (Nothing, Nothing)
			Just (Left v)	-> (Just v, Nothing)
			Just (Right v)	-> (Nothing, Just v)
		ts' = mapWithKeyAlg (\ x -> Elem . mapEitherEdge (\ xs -> f (ks ++ x:xs))) ts
		tsL = mapMaybeAlg (\ _ (Elem (tsL, _)) -> tsL) ts'
		tsR = mapMaybeAlg (\ _ (Elem (_, tsR)) -> tsR) ts'

mapAppEdge :: (Applicative f, TrieKey k m, Sized w) => ([k] -> v -> f w) -> Edge k m v -> f (Edge k m w)
mapAppEdge f (Edge _ ks v ts) = liftA2 (edge ks) (traverse (f ks) v) (mapAppAlg (\ x -> mapAppEdge (\ xs -> f (ks ++ x:xs))) ts)

unionMaybeEdge :: (Eq k, TrieKey k m, Sized v) => ([k] -> v -> v -> Maybe v) -> Edge k m v -> Edge k m v -> MEdge k m v
unionMaybeEdge f (Edge nK ks0 vK tsK) (Edge nL ls0 vL tsL) = procEdge 0 ks0 ls0 where
	procEdge i _ _ | i `seq` False = undefined
	procEdge i (k:ks) (l:ls)
		| k == l	= procEdge (i+1) ks ls
		| otherwise	= Just $ Edge (nK + nL) (take i ks0) Nothing $ fromListAlg' [(k, Edge nK ks vK tsK), (l, Edge nL ls vL tsL)]
	procEdge _ [] (l:ls) = guardNullEdge $ edge ks0 vK $ alterAlg g l tsK
	  where	g Nothing = Just (Edge nL ls vL tsL)
		g (Just e') = unionMaybeEdge (\ ls' -> f (ks0 ++ l:ls')) e' (Edge nL ls vL tsL)
	procEdge _ (k:ks) [] = guardNullEdge $ edge ls0 vL $ alterAlg g k tsL 
	  where	g Nothing = Just $ Edge nK ks vK tsK
		g (Just e') = unionMaybeEdge (\ ks' -> f (ls0 ++ k:ks')) (Edge nK ks vK tsK) e'
	procEdge _ [] [] = guardNullEdge $ edge ks0 (unionMaybe (f ks0) vK vL) $
		unionMaybeAlg (\ x -> unionMaybeEdge (\ xs -> f (ks0 ++ x:xs))) tsK tsL

intersectEdge :: (Eq k, TrieKey k m, Sized c) => ([k] -> a -> b -> Maybe c) -> Edge k m a -> Edge k m b -> MEdge k m c
intersectEdge f (Edge nK ks0 vK tsK) (Edge nL ls0 vL tsL) = procEdge ks0 ls0 where
	procEdge (k:ks) (l:ls)
		| k == l	= procEdge ks ls
		| otherwise	= Nothing
	procEdge (k:ks) [] = do
		e' <- lookupAlg k tsL
		Edge nX xs vX tsX <- intersectEdge (\ ks' -> f (ls0 ++ k:ks')) (Edge nK ks vK tsK) e'
		return (Edge nX (ls0 ++ k:xs) vX tsX)
	procEdge [] (l:ls) = do
		e' <- lookupAlg l tsK
		Edge nX xs vX tsX <- intersectEdge (\ ls' -> f (ks0 ++ l:ls')) e' (Edge nL ls vL tsL)
		return (Edge nX (ks0 ++ l:xs) vX tsX)
	procEdge [] [] = guardNullEdge $ edge ks0 (intersectMaybe (f ks0) vK vL) $
		intersectAlg (\ x -> intersectEdge (\ xs -> f (ks0 ++ x:xs))) tsK tsL

{-# SPECIALIZE differenceEdge :: (Eq k, TrieKey k m) => ([k] -> Elem v -> w -> Maybe (Elem v)) -> 
	Edge k m (Elem v) -> Edge k m w -> MEdge k m (Elem v) #-}
differenceEdge :: (Eq k, TrieKey k m, Sized v) => ([k] -> v -> w -> Maybe v) -> Edge k m v -> Edge k m w -> MEdge k m v
differenceEdge f e@(Edge nK ks0 vK tsK) (Edge nL ls0 vL tsL) = procEdge ks0 ls0 where
	procEdge (k:ks) (l:ls)
		| k == l	= procEdge ks ls
	procEdge (k:ks) []
		| Just e' <- lookupAlg k tsL
		= do	Edge nX xs vX tsX <- differenceEdge (\ ks' -> f (ls0 ++ k:ks')) (Edge nK ks vK tsK) e'
			return (Edge nX (ls0 ++ k:xs) vX tsX)
	procEdge [] (l:ls) = guardNullEdge $ edge ks0 vK $ alterAlg g l tsK
	  where	g Nothing = Nothing
	  	g (Just e') = differenceEdge (\ ls' -> f (ks0 ++ l:ls')) e' (Edge nL ls vL tsL)
	procEdge [] [] = guardNullEdge $ edge ks0 (differenceMaybe (f ks0) vK vL) $
		differenceAlg (\ x -> differenceEdge (\ xs -> f (ks0 ++ x:xs))) tsK tsL
	procEdge _ _ = Just e

{-# SPECIALIZE getMinEdge :: TrieKey k m => Edge k m (Elem v) -> (([k], Elem v), MEdge k m (Elem v)) #-}
getMinEdge :: (Sized v, TrieKey k m) => Edge k m v -> (([k], v), MEdge k m v)
getMinEdge (Edge nK ks (Just v) ts) = ((ks, v), guardNullEdge $ Edge (nK - getSize v) ks Nothing ts)
getMinEdge (Edge nK ks _ ts) 
	| Just ((l, e), ts') <- getMinAlg ts, ((ls, v), e') <- getMinEdge e
		= ((ks ++ l:ls, v), guardNullEdge $ edge ks Nothing $ maybe ts' (\ e' -> snd $ updateMinAlg (\ _ _ -> (False, Just e')) ts) e')
getMinEdge _ = error "Uncompacted edge"

getMaxEdge :: (Sized v, TrieKey k m) => Edge k m v -> (([k], v), MEdge k m v)
getMaxEdge (Edge nK ks v0 ts)
	| Just ((l, e), ts') <- getMaxAlg ts, ((ls, v), e') <- getMaxEdge e
		= ((ks ++ l:ls, v), guardNullEdge $ edge ks v0 $ maybe ts' (\ e' -> snd $ updateMaxAlg (\ _ _ -> (False, Just e')) ts) e')
getMaxEdge (Edge nK ks (Just v) ts) = ((ks, v), guardNullEdge $ Edge (nK - getSize v) ks Nothing ts)
getMaxEdge _ = error "Uncompacted edge"

updateMinEdge :: (TrieKey k m, Sized v) => ([k] -> v -> (Bool, Maybe v)) -> Edge k m v -> (Bool, MEdge k m v)
updateMinEdge f (Edge _ ks (Just v) ts) = fmap (\ v' -> guardNullEdge $ edge ks v' ts) (f ks v)
updateMinEdge f (Edge _ ks Nothing ts)
	= fmap (guardNullEdge . edge ks Nothing) $ updateMinAlg (\ l -> updateMinEdge (\ ls -> f (ks ++ l:ls))) ts

updateMaxEdge :: (TrieKey k m, Sized v) => ([k] -> v -> (Bool, Maybe v)) -> Edge k m v -> (Bool, MEdge k m v)
updateMaxEdge f (Edge _ ks (Just v) ts)
	| nullAlg ts	= fmap (\ v' -> guardNullEdge $ edge ks v' ts) (f ks v)
updateMaxEdge f (Edge _ ks v ts) = 
	fmap (guardNullEdge . edge ks v) $ updateMaxAlg (\ l -> updateMaxEdge (\ ls -> f (ks ++ l:ls))) ts

isSubmapEdge :: TrieKey k m => (k -> k -> Bool) -> (a -> b -> Bool) -> (k -> m (Edge k m b) -> MEdge k m b) -> (m (Edge k m a) -> m (Edge k m b) -> Bool) -> 
	Edge k m a -> Edge k m b -> Bool
isSubmapEdge (==) (<=) lookup (<<=) (Edge nK ks vK tsK) (Edge nL ls vL tsL) = procEdge ks ls where
	procEdge (k:ks) (l:ls)
		| k == l	= procEdge ks ls
	procEdge (k:ks) []
		| Just e <- lookup k tsL
			= isSubmapEdge (==) (<=) lookup (<<=) (Edge nK ks vK tsK) e
	procEdge [] [] 
		| Nothing <- vK	= tsK <<= tsL
		| Just x <- vK, Just y <- vL, x <= y
				= tsK <<= tsL
	procEdge _ _ = False
validEdge :: TrieKey k m => Edge k m v -> Bool
validEdge (Edge _ _ Nothing m)
	| nullAlg m	= False
	| Just{} <- getSingleAlg m
			= False
validEdge (Edge _ _ _ m)
	= valid m && all validEdge m


splitEdge :: (Ord k, TrieKey k m, Sized a) => (a -> (Maybe a, Maybe b, Maybe a)) -> [k] -> Edge k m a -> (MEdge k m a, Maybe b, MEdge k m a)
splitEdge f ks0 e@(Edge nL ls0 v ts) = procEdge ks0 ls0 where
	answerLess = (Nothing, Nothing, Just e) -- if ks0 < ls0
	answerMore = (Just e, Nothing, Nothing) -- if ks0 > ls0 
	procEdge (k:ks) (l:ls) = case compare k l of
		LT	-> answerLess
		EQ	-> procEdge ks ls
		GT	-> answerMore
	procEdge (k:ks) [] = case splitLookupAlg (splitEdge f ks) k ts of
		(tsL, ans, tsR)	-> (guardNullEdge $ edge ls0 Nothing tsL, ans, guardNullEdge $ edge ls0 v tsR)
	procEdge [] (l:ls) = answerLess
	procEdge [] [] 
		| Just v <- v, (vL, ans, vR) <- f v
			= (fmap (\ v' -> edge ls0 (Just v') emptyAlg) vL, ans, 
				guardNullEdge $ edge ls0 vR ts)
		| otherwise	= answerLess -- all children of e match ks0 initially but are longer, and v is Nothing