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

{- |
This module implements the functionality of a /monoid queue/, which is essentially a priority queue that merges values with equal keys.  Several implementations were considered:

* A pairing heap which, as part of its merging operation, merges any identical-keyed values it encounters.  This may result in partial merging of equal-keyed values for several different keys
	during a single delete-min operation, decreasing the number of nodes in the queue without costing any additional comparisons.  In addition, it naturally falls out that
	the values associated with the minimum key are always fully merged.  Disadvantages include no control over the balancedness of complete merges, and possible extra polymorphism overhead.
* This considerably simpler bootstrap on a vanilla @PQueue@, which keeps the totally merged value associated with its very minimum key and performs no partial merging until a
	key becomes the minimum. (The fact that no partial merging is performed allows an optimized and balanced 'mconcat' to be used on all the values associated with a key at once.)

The primary difference between these two is the compromise between /performance of the heap itself/ and /performance of the merging of the monoids/:
	* The first implementation speculatively takes advantage of opportunities to merge nodes with equal keys.  No actual extra work is being done -- the same number of comparisons
		is made, and the actual merging is performed lazily.  As a result, truly the only disadvantages of this approach are the completely uncontrolled balance of the merges,
		and possible polymorphism overhead that can be carefully hand-removed.
	* The second implementation allows an optimized bulk merge operation, which -- in the case of tries, which is what after all the motivation for this structure --
		has extremely significant advantages.

A version of the second implementation is included with the Cabal distribution in 
"Data.Queue.TrieQueue.MonoidQueue2", and can serve as a literal drop-in replacement for this module.
See its implementation notes for further details.

-}
module Data.Queue.TrieQueue.MonoidQueue (MQueue, extractSingle, replace) where

import Data.Queue.Class
import Data.Queue.QueueHelpers(fusing')
import Data.Maybe
import Data.Monoid
import qualified Data.Tree as Tree

import Data.Queue.TrieQueue.Edge

-- | A pairing heap node in the monoid queue; nonempty.
data MNode k m = Node k m (MForest k m)
type MForest k m = [MNode k m]
-- | A full-fledged priority queue, including empty queues.
newtype MQueue k m = MQ (Maybe (MNode k m)) deriving (Show)

instance (Show k, Show m) => Show (MNode k m) where
	show = Tree.drawTree . Tree.unfoldTree (\ (Node k m ts) -> (show (k :- m), ts))

-- | A @Functor@ instance exploited for great justice in @mergeNodes@.
instance Functor (MNode k) where
	fmap f (Node k m ns) = Node k (f m) (map (fmap f) ns)

instance (Ord k, Monoid m) => Monoid (MNode k m) where
	-- no mempty declaration
	mappend = mergeMNode $! mappend

instance (Ord k, Monoid m) => Monoid (MQueue k m) where
	mempty = MQ Nothing
	MQ n1 `mappend` MQ n2 = MQ (n1 `mappend` n2)
	 -- This is the straightforward implementation:
	 -- mconcat qs = MQ $ fusing [n | MQ (Just n) <- qs]
	 -- This implementation allows optimized mass mconcats rather than just using individual mappends,
	 -- and exploits the flexible monoid structure.
	mconcat qs = mergeNodes [n | MQ (Just n) <- qs]


mergeMNode :: Ord k => (m -> m -> m) -> MNode k m -> MNode k m -> MNode k m
mergeMNode (><) n1@(Node k1 m1 ns1) n2@(Node k2 m2 ns2) = case compare k1 k2 of
	LT	-> Node k1 m1 (n2:ns1)
	EQ	-> Node k1 (m1 >< m2) (ns1 ++ ns2)
	GT	-> Node k2 m2 (n1:ns2)

{-# INLINE mergeNodes' #-}
-- | Merges a collection of nodes by performing a balanced fuse and performing an mconcat on all blocks of equal-keyed monoid values.
mergeNodes' :: (Ord k) => (k -> k -> Ordering) -> (m -> m -> m) -> MForest k m -> MQueue k m
--mergeNodes ns = MQ (fmap (fmap mconcat) $ fusing (map (fmap $ \ m -> [m]) ns))
mergeNodes' _ (><) = MQ . fusing' (mergeMNode (><))

{-# INLINE mergeNodes #-}
mergeNodes :: (Ord k, Monoid m) => MForest k m -> MQueue k m
mergeNodes = (mergeNodes' $! compare) $! mappend

instance (Ord k, Monoid m) => IQueue (MQueue k m) where
	type QueueKey (MQueue k m) = k :- m
	empty = mempty
	merge = mappend
	mergeAll = mconcat

	singleton (k :- m) = MQ (Just (Node k m []))
	insertAll = merge . fromListMQ
	fromList = fromListMQ
	extract = extractMQ

	null (MQ Nothing) = True
	null _ = False

{-# INLINE single #-}
single :: k :- m -> MNode k m
single (k :- m) = Node k m []

extractMQ :: (Ord k, Monoid m) => MQueue k m -> Maybe (k :- m, MQueue k m)
extractMQ (MQ t) = fmap extract' t where
	extract' (Node k m ns) = (k :- m, mergeNodes ns)

fromListMQ :: (Ord k, Monoid m) => [k :- m] -> MQueue k m
fromListMQ ks = mergeNodes $ map single ks

{-# INLINE extractSingle #-}
extractSingle :: (Ord k, Monoid m) => MQueue k m -> Maybe (k :- m)
extractSingle (MQ (Just (Node k m []))) = Just (k :- m)
extractSingle _ = Nothing

-- UNSAFE.  Do not use unless you know the key you're putting in is less than or equal to the minimum key.
replace :: (Ord k, Monoid m) => (k :- m) -> MQueue k m -> MQueue k m
replace km@(k :- m) (MQ t) = MQ $ Just (maybe (single km) (\ (Node _ _ ns) -> Node k m ns) t)