{-# LANGUAGE PatternGuards, TypeFamilies, GeneralizedNewtypeDeriving #-}
{-# OPTIONS -fno-warn-missing-methods #-}

-- | An experimental trie-based priority queue for lists.
module Data.Queue.TrieQueue where
-- Data.Sequence-labelled trie implementation, bootstrapping monoid structure to achieve maximum great justice.

import Data.Queue.Class
import Data.Queue.QueueHelpers

import Control.Arrow((***))
import Control.Monad(liftM2)

import Data.Function
import Data.Maybe
import Data.List (sortBy, groupBy)
import Data.Monoid
import Data.Ord
import Data.Sequence (Seq, viewl, ViewL(..), (><), (|>), (<|))
import qualified Data.Sequence as Seq
import Data.Map (Map, findMin, minViewWithKey, fromDistinctAscList)
import qualified Data.Map as Map
import qualified Data.Foldable as Fold (toList)

import GHC.Exts

type Label e = Seq e

-- Type of a nonempty trie.
data Trie e = 	Leaf (Label e) {-# UNPACK #-} !Int 	-- Leaf xs n represents n occurrences of the string xs.  n is always strictly positive.
	      | Edge (Label e) {-# UNPACK #-} !Int (Map e (Trie e)) -- Edge xs n m represents n occurrences of xs, and xs prepended to each element of the map.
		deriving (Eq)
type MTrie e = Maybe (Trie e)
newtype TrieQueue e = TQ (HeapQ (Trie e)) deriving (Eq, Show, Monoid)

{-# INLINE mkTQ #-}
mkTQ :: Int -> Trie e -> TrieQueue e
mkTQ n t = TQ (HQ n (Just t))

mkLab :: [e] -> Label e
mkLab = Seq.fromList

instance Show e => Show (Trie e) where
	show (Leaf xs xn) = "(" ++ show xn ++ "x" ++ show (Fold.toList xs) ++ ")"
	show (Edge xs xn m) = "==" ++ show xn ++ "x" ++ show (Fold.toList xs) ++ "==>" ++ show m

instance Ord e => Monoid (Trie e) where
	-- not a true monoid instance; only a semigroup
	mappend = merger

instance Ord e => IQueue (TrieQueue e) where
	type QueueKey (TrieQueue e) = [e]
	empty = mempty
	merge = mappend
	mergeAll = mconcat

	singleton xs = mkTQ 1 $ Leaf (mkLab xs) 1
--	fromList = TQ . liftM2 HQ length trieFromList

	extract (TQ (HQ n t)) = fmap ((Fold.toList *** (TQ . HQ (n-1))) . extractMin') t
	size (TQ (HQ n _)) = n

catTrie :: Ord e => Label e -> Trie e -> Trie e
catTrie xs (Leaf ys n) = Leaf (xs >< ys) n
catTrie xs (Edge ys n m) = compactTrie (Edge (xs >< ys) n m)

compactTrie :: Ord e => Trie e -> Trie e
compactTrie = fromJust . compactTrie'

compactTrie' :: Ord e => Trie e -> MTrie e
compactTrie' t@(Edge xs 0 m)
	| Map.null m	= Nothing
	| Map.size m == 1, (y, yT) <- findMin m
			= Just $ catTrie (xs |> y) yT
	| otherwise	= Just t
compactTrie' (Edge xs n m)
	| Map.null m	= Just $ Leaf xs n
compactTrie' (Leaf _ 0) = Nothing
compactTrie' t = Just t


extractMin' :: Ord e => Trie e -> (Label e, MTrie e)
extractMin' (Leaf xs n) = (xs, compactTrie' (Leaf xs (n-1)))
extractMin' (Edge xs (n+1) m) = (xs, compactTrie' (Edge xs n m))
extractMin' (Edge xs 0 m)
	| Just ((y, yT), m') <- minViewWithKey m,
		(ys, yT') <- extractMin' yT
	= (xs >< (y <| ys), maybe (compactTrie' (Edge xs 0 m')) (\ yT' -> Just $ Edge xs 0 $ Map.insert y yT' m') yT')
extractMin' _ = error "Internal failure to note empty queue"

--extractMin :: Ord e => Trie e -> Maybe ([e], Trie e)
--extractMin = fmap (first toList) . extractMin'

type TailMaker e = Label e -> Trie e

{-# INLINE merge' #-}
merge' :: Ord e => Label e -> Label e -> TailMaker e -> TailMaker e -> (e -> TailMaker e) -> (e -> TailMaker e) -> Trie e -> Trie e
merge' xs0 ys0 xTail yTail xCons yCons xy = merge'' 0 xs0 ys0 where
	merge'' n xs ys = case (viewl xs, viewl ys) of
		(x :< xs1, y :< ys1) -> let pfx = Seq.take n xs0; xT = xTail xs1; yT = yTail ys1; in case x `compare` y of
			LT	-> Edge pfx 0 $ fromDistinctAscList [(x, xT), (y, yT)]
			EQ	-> merge'' (n+1) xs1 ys1
			GT	-> Edge pfx 0 $ fromDistinctAscList [(y, yT), (x, xT)]
		(x :< xs1, EmptyL)	-> yCons x xs1
		(EmptyL, y :< ys1)	-> xCons y ys1
		(EmptyL, EmptyL)	-> xy
	
merger :: Ord e => Trie e -> Trie e -> Trie e	
Leaf xs0 xn `merger` Leaf ys0 yn =
	merge' xs0 ys0 (flip Leaf xn) (flip Leaf yn) (edger xs0 xn yn) (edger ys0 yn xn) (Leaf xs0 (xn + yn))
	where	edger xs xn yn y ys = Edge xs xn $ Map.singleton y (Leaf ys yn)
Leaf xs0 xn `merger` Edge ys0 yn yM =
	merge' xs0 ys0 (flip Leaf xn) (\ ys -> Edge ys yn yM) (\ y ys -> Edge xs0 xn $ Map.singleton y (Edge ys yn yM))
		(\ x xs -> Edge ys0 yn $ Map.insertWith merger x (Leaf xs xn) yM) (Edge xs0 (xn + yn) yM)
x@Edge{} `merger` y@Leaf{} = merger y x
Edge xs0 xn xM `merger` Edge ys0 yn yM
	= merge' xs0 ys0 (edger xn xM) (edger yn yM) (cons xs0 xn yn yM) (cons ys0 yn xn xM)
			(Edge xs0 (xn + yn) $ Map.unionWith merger xM yM)
	where	edger n m l = Edge l n m
		cons xs xn yn yM y = Edge xs xn . Map.singleton y . edger yn yM

trieFromList :: Ord e => [[e]] -> MTrie e
trieFromList = extractCommon mempty
	where	groupHeads = groupBy ((==) `on` listToMaybe) . sortBy (comparing listToMaybe)
		extractCommon pfx xs = case groupHeads xs of
			[]			-> Nothing
			[empties@([]:_)]	-> Just $ Leaf pfx (length empties)
			(empties@([]:_):xss)	-> Just $ Edge pfx (length empties) (fromGroups xss)
				-- even if there's only one other group, we end the edge here
				-- a more optimized implementation might specialize for this case
			[(y:ys):yss]		-> extractCommon (pfx |> y) (ys:map tail yss)
				-- if there's but a single group with a shared first character, snoc it onto the accumulated prefix
				-- and recurse
			xss			-> Just $ Edge pfx 0 (fromGroups xss)
		fromGroups xss = fromDistinctAscList [(y, fromJust $ trieFromList $ ys : map tail yss) | ((y:ys):yss) <- xss]