{-# LANGUAGE CPP #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Data.PQueue.Prio.Max
-- Copyright   :  (c) Louis Wasserman 2010
-- License     :  BSD-style
-- Maintainer  :  libraries@haskell.org
-- Stability   :  experimental
-- Portability :  portable
--
-- General purpose priority queue, supporting extract-minimum operations.
-- Each element is associated with a /key/, and the priority queue supports
-- viewing and extracting the element with the minimum key.
--
-- An amortized running time is given for each operation, with /n/ referring
-- to the length of the sequence and /i/ being the integral index used by
-- some operations.  These bounds hold even in a persistent (shared) setting.
--
-- This implementation is based on a binomial heap augmented with a global root.
-- The spine of the heap is maintained lazily.
--
-- This implementation does not guarantee stable behavior.  Ties are broken
-- arbitrarily -- that is, if @k1 <= k2@ and @k2 <= k1@, then there are no
-- guarantees about the relative order in which @k1@, @k2@, and their associated
-- elements are returned.
-- 
-- This implementation offers a number of methods of the form @xxxU@, where @U@ stands for
-- "unordered."  No guarantees are made on the execution or traversal order of
-- these functions.
-----------------------------------------------------------------------------
module Data.PQueue.Prio.Max (
	MaxPQueue,
	-- * Construction
	empty,
	singleton,
	insert,
	union,
	unions, 
	-- * Query
	null,
	size,
	-- ** Maximum view
	findMax,
	getMax,
	deleteMax,
	deleteFindMax,
	alterMax,
	alterMaxWithKey,
	updateMax,
	updateMaxWithKey,
	maxView,
	maxViewWithKey,
	-- * Traversal
	-- ** Map
	map,
	mapWithKey,
	mapKeys,
	mapKeysMonotonic,
	-- ** Fold
	foldrWithKey,
	foldlWithKey,
	-- ** Traverse
	traverseWithKey,
	-- * Subsets
	-- ** Indexed
	take,
	drop,
	splitAt,
	-- ** Predicates
	takeWhile,
	takeWhileWithKey,
	dropWhile,
	dropWhileWithKey,
	span,
	spanWithKey,
	break,
	breakWithKey,
	-- *** Filter
	filter,
	filterWithKey,
	partition,
	partitionWithKey,
	mapMaybe,
	mapMaybeWithKey,
	mapEither,
	mapEitherWithKey,
	-- * List operations
	-- ** Conversion from lists
	fromList,
	fromAscList,
	fromDescList,
	-- ** Conversion to lists
	keys,
	elems,
	assocs,
	toAscList,
	toDescList,
	toList,
	-- * Unordered operations
	foldrU,
	foldrWithKeyU,
	foldlU,
	foldlWithKeyU,
	traverseU,
	traverseWithKeyU,
	keysU,
	elemsU,
	assocsU,
	toListU,
	-- * Helper methods
	seqSpine
	)
	where

import Control.Applicative hiding (empty)
import Control.Arrow
import Data.Monoid
import qualified Data.List as List
import Data.Foldable hiding (toList)
import Data.Traversable
import Data.Maybe hiding (mapMaybe)

import Prelude hiding (map, filter, break, span, takeWhile, dropWhile, splitAt, take, drop, (!!), null, foldr, foldl)

import qualified Data.PQueue.Prio.Min as Q

#ifdef __GLASGOW_HASKELL__
import GHC.Exts (build)
import Text.Read (Lexeme(Ident), lexP, parens, prec,
	readPrec, readListPrec, readListPrecDefault)
import Data.Data
#else
build :: ((a -> [a] -> [a]) -> [a] -> [a]) -> [a]
build f = f (:) []
#endif

first' :: (a -> b) -> (a, c) -> (b, c)
first' f (a, c) = (f a, c)

second' :: (b -> c) -> (a, b) -> (a, c)
second' f (a, b) = (a, f b)

newtype Down a = Down {unDown :: a} deriving (Eq)

-- | A priority queue where values of type @a@ are annotated with keys of type @k@.
-- The queue supports extracting the element with maximum key.
newtype MaxPQueue k a = MaxPQ (Q.MinPQueue (Down k) a) deriving (Eq, Ord)

instance Ord a => Ord (Down a) where
	Down a `compare` Down b = b `compare` a
	Down a <= Down b = b <= a

instance Functor Down where
	fmap f (Down a) = Down (f a)

instance Functor (MaxPQueue k) where
	fmap f (MaxPQ q) = MaxPQ (fmap f q)

instance Ord k => Foldable (MaxPQueue k) where
	foldr f z (MaxPQ q) = foldr f z q
	foldl f z (MaxPQ q) = foldl f z q

instance Ord k => Traversable (MaxPQueue k) where
	traverse f (MaxPQ q) = MaxPQ <$> traverse f q

-- | /O(1)/.  Returns the empty priority queue.
empty :: MaxPQueue k a
empty = MaxPQ Q.empty

-- | /O(1)/.  Constructs a singleton priority queue.
singleton :: k -> a -> MaxPQueue k a
singleton k a = MaxPQ (Q.singleton (Down k) a)

-- | Amortized /O(1)/, worst-case /O(log n)/.  Inserts
-- an element with the specified key into the queue.
insert :: Ord k => k -> a -> MaxPQueue k a -> MaxPQueue k a
insert k a (MaxPQ q) = MaxPQ (Q.insert (Down k) a q)

-- | Amortized /O(log(min(n1, n2)))/, worst-case /O(log(max(n1, n2)))/.  Returns the union
-- of the two specified queues.
union :: Ord k => MaxPQueue k a -> MaxPQueue k a -> MaxPQueue k a
MaxPQ q1 `union` MaxPQ q2 = MaxPQ (q1 `Q.union` q2)

-- | The union of a list of queues: (@'unions' == 'List.foldl' 'union' 'empty'@).
unions :: Ord k => [MaxPQueue k a] -> MaxPQueue k a
unions qs = MaxPQ (Q.unions [q | MaxPQ q <- qs])

-- | /O(1)/.  Checks if this priority queue is empty.
null :: MaxPQueue k a -> Bool
null (MaxPQ q) = Q.null q

-- | /O(1)/.  Returns the size of this priority queue.
size :: MaxPQueue k a -> Int
size (MaxPQ q) = Q.size q

-- | /O(1)/.  The maximal (key, element) in the queue.  Calls 'error' if empty.
findMax :: MaxPQueue k a -> (k, a)
findMax = fromMaybe (error "Error: findMax called on an empty queue") . getMax

-- | /O(1)/.  The maximal (key, element) in the queue, if the queue is nonempty.
getMax :: MaxPQueue k a -> Maybe (k, a)
getMax (MaxPQ q) = do
	(Down k, a) <- Q.getMin q
	return (k, a)

-- | /O(log n)/.  Delete and find the element with the maximum key.  Calls 'error' if empty.
deleteMax :: Ord k => MaxPQueue k a -> MaxPQueue k a
deleteMax (MaxPQ q) = MaxPQ (Q.deleteMin q)

-- | /O(log n)/.  Delete and find the element with the maximum key.  Calls 'error' if empty.
deleteFindMax :: Ord k => MaxPQueue k a -> ((k, a), MaxPQueue k a)
deleteFindMax = fromMaybe (error "Error: deleteFindMax called on an empty queue") . maxViewWithKey

-- | /O(1)/.  Alter the value at the maximum key.  If the queue is empty, does nothing.
alterMax :: (a -> a) -> MaxPQueue k a -> MaxPQueue k a 
alterMax = alterMaxWithKey . const

-- | /O(1)/.  Alter the value at the maximum key.  If the queue is empty, does nothing.
alterMaxWithKey :: (k -> a -> a) -> MaxPQueue k a -> MaxPQueue k a
alterMaxWithKey f (MaxPQ q) = MaxPQ (Q.alterMinWithKey (f . unDown) q)

-- | /O(log n)/.  (Actually /O(1)/ if there's no deletion.)  Update the value at the maximum key.
-- If the queue is empty, does nothing.
updateMax :: Ord k => (a -> Maybe a) -> MaxPQueue k a -> MaxPQueue k a
updateMax = updateMaxWithKey . const

-- | /O(log n)/.  (Actually /O(1)/ if there's no deletion.)  Update the value at the maximum key.
-- If the queue is empty, does nothing.
updateMaxWithKey :: Ord k => (k -> a -> Maybe a) -> MaxPQueue k a -> MaxPQueue k a
updateMaxWithKey f (MaxPQ q) = MaxPQ (Q.updateMinWithKey (f . unDown) q)

-- | /O(log n)/.  Retrieves the value associated with the maximum key of the queue, and the queue
-- stripped of that element, or 'Nothing' if passed an empty queue.
maxView :: Ord k => MaxPQueue k a -> Maybe (a, MaxPQueue k a)
maxView q = do
	((_, a), q') <- maxViewWithKey q
	return (a, q')

-- | /O(log n)/.  Retrieves the maximal (key, value) pair of the map, and the map stripped of that
-- element, or 'Nothing' if passed an empty map.
maxViewWithKey :: Ord k => MaxPQueue k a -> Maybe ((k, a), MaxPQueue k a)
maxViewWithKey (MaxPQ q) = do
	((Down k, a), q') <- Q.minViewWithKey q
	return ((k, a), MaxPQ q')

-- | /O(n)/.  Map a function over all values in the queue.
map :: (a -> b) -> MaxPQueue k a -> MaxPQueue k b
map = mapWithKey . const

-- | /O(n)/.  Map a function over all values in the queue.
mapWithKey :: (k -> a -> b) -> MaxPQueue k a -> MaxPQueue k b
mapWithKey f (MaxPQ q) = MaxPQ (Q.mapWithKey (f . unDown) q)

-- | /O(n)/.  Map a function over all values in the queue.
mapKeys :: Ord k' => (k -> k') -> MaxPQueue k a -> MaxPQueue k' a
mapKeys f (MaxPQ q) = MaxPQ (Q.mapKeys (fmap f) q)

-- | /O(n)/.  @'mapKeysMonotonic' f q == 'mapKeys' f q@, but only works when @f@ is strictly
-- monotonic.  /The precondition is not checked./  This function has better performance than
-- 'mapKeys'.
mapKeysMonotonic :: (k -> k') -> MaxPQueue k a -> MaxPQueue k' a
mapKeysMonotonic f (MaxPQ q) = MaxPQ (Q.mapKeysMonotonic (fmap f) q)

-- | /O(n log n)/.  Fold the keys and values in the map, such that 
-- @'foldrWithKey' f z q == 'List.foldr' ('uncurry' f) z ('toAscList' q)@.
-- 
-- If you do not care about the traversal order, consider using 'foldrWithKeyU'.
foldrWithKey :: Ord k => (k -> a -> b -> b) -> b -> MaxPQueue k a -> b
foldrWithKey f z (MaxPQ q) = Q.foldrWithKey (f . unDown) z q

-- | /O(n log n)/.  Fold the keys and values in the map, such that 
-- @'foldlWithKey' f z q == 'List.foldl' ('uncurry' . f) z ('toAscList' q)@.
-- 
-- If you do not care about the traversal order, consider using 'foldlWithKeyU'.
foldlWithKey :: Ord k => (b -> k -> a -> b) -> b -> MaxPQueue k a -> b
foldlWithKey f z (MaxPQ q) = Q.foldlWithKey (\ z -> f z . unDown) z q

-- | /O(n log n)/.  Traverses the elements of the queue in descending order by key.
-- (@'traverseWithKey' f q == 'fromDescList' <$> 'traverse' ('uncurry' f) ('toDescList' q)@)
-- 
-- If you do not care about the /order/ of the traversal, consider using 'traverseWithKeyU'.
traverseWithKey :: (Ord k, Applicative f) => (k -> a -> f b) -> MaxPQueue k a -> f (MaxPQueue k b)
traverseWithKey f (MaxPQ q) = MaxPQ <$> Q.traverseWithKey (f . unDown) q

-- | /O(k log n)/.  Takes the first @k@ (key, value) pairs in the queue, or the first @n@ if @k >= n@.
-- (@'take' k q == 'List.take' k ('toDescList' q)@)
take :: Ord k => Int -> MaxPQueue k a -> [(k, a)]
take k (MaxPQ q) = fmap (first' unDown) (Q.take k q)

-- | /O(k log n)/.  Deletes the first @k@ (key, value) pairs in the queue, or returns an empty queue if @k >= n@.
drop :: Ord k => Int -> MaxPQueue k a -> MaxPQueue k a
drop k (MaxPQ q) = MaxPQ (Q.drop k q)

-- | /O(k log n)/.  Equivalent to @('take' k q, 'drop' k q)@.
splitAt :: Ord k => Int -> MaxPQueue k a -> ([(k, a)], MaxPQueue k a)
splitAt k (MaxPQ q) = case Q.splitAt k q of
	(xs, q') -> (fmap (first' unDown) xs, MaxPQ q')

-- | Takes the longest possible prefix of elements satisfying the predicate.
-- (@'takeWhile' p q == 'List.takeWhile' (p . 'snd') ('toAscList' q)@)
takeWhile :: Ord k => (a -> Bool) -> MaxPQueue k a -> [(k, a)]
takeWhile = takeWhileWithKey . const

-- | Takes the longest possible prefix of elements satisfying the predicate.
-- (@'takeWhile' p q == 'List.takeWhile' (uncurry p) ('toAscList' q)@)
takeWhileWithKey :: Ord k => (k -> a -> Bool) -> MaxPQueue k a -> [(k, a)]
takeWhileWithKey p (MaxPQ q) = fmap (first' unDown) (Q.takeWhileWithKey (p . unDown) q)

-- | Removes the longest possible prefix of elements satisfying the predicate.
dropWhile :: Ord k => (a -> Bool) -> MaxPQueue k a -> MaxPQueue k a
dropWhile = dropWhileWithKey . const

-- | Removes the longest possible prefix of elements satisfying the predicate.
dropWhileWithKey :: Ord k => (k -> a -> Bool) -> MaxPQueue k a -> MaxPQueue k a
dropWhileWithKey p (MaxPQ q) = MaxPQ (Q.dropWhileWithKey (p . unDown) q)

-- | Equivalent to @('takeWhile' p q, 'dropWhile' p q)@.
span :: Ord k => (a -> Bool) -> MaxPQueue k a -> ([(k, a)], MaxPQueue k a)
span = spanWithKey . const

-- | Equivalent to @'span' ('not' . p)@.
break :: Ord k => (a -> Bool) -> MaxPQueue k a -> ([(k, a)], MaxPQueue k a)
break = breakWithKey . const

-- | Equivalent to @'spanWithKey' (\ k a -> 'not' (p k a)) q@.
spanWithKey :: Ord k => (k -> a -> Bool) -> MaxPQueue k a -> ([(k, a)], MaxPQueue k a)
spanWithKey p (MaxPQ q) = case Q.spanWithKey (p . unDown) q of
	(xs, q') -> (fmap (first' unDown) xs, MaxPQ q')

-- | Equivalent to @'spanWithKey' (\ k a -> 'not' (p k a)) q@.
breakWithKey :: Ord k => (k -> a -> Bool) -> MaxPQueue k a -> ([(k, a)], MaxPQueue k a)
breakWithKey p (MaxPQ q) = case Q.breakWithKey (p . unDown) q of
	(xs, q') -> (fmap (first' unDown) xs, MaxPQ q')

-- | /O(n)/.  Filter all values that satisfy the predicate.
filter :: Ord k => (a -> Bool) -> MaxPQueue k a -> MaxPQueue k a
filter = filterWithKey . const

-- | /O(n)/.  Filter all values that satisfy the predicate.
filterWithKey :: Ord k => (k -> a -> Bool) -> MaxPQueue k a -> MaxPQueue k a
filterWithKey p (MaxPQ q) = MaxPQ (Q.filterWithKey (p . unDown) q)

-- | /O(n)/.  Partition the queue according to a predicate.  The first queue contains all elements
-- which satisfy the predicate, the second all elements that fail the predicate.
partition :: Ord k => (a -> Bool) -> MaxPQueue k a -> (MaxPQueue k a, MaxPQueue k a)
partition = partitionWithKey . const

-- | /O(n)/.  Partition the queue according to a predicate.  The first queue contains all elements
-- which satisfy the predicate, the second all elements that fail the predicate.
partitionWithKey :: Ord k => (k -> a -> Bool) -> MaxPQueue k a -> (MaxPQueue k a, MaxPQueue k a)
partitionWithKey p (MaxPQ q) = case Q.partitionWithKey (p . unDown) q of
	(q1, q0) -> (MaxPQ q1, MaxPQ q0)

-- | /O(n)/.  Map values and collect the 'Just' results.
mapMaybe :: Ord k => (a -> Maybe b) -> MaxPQueue k a -> MaxPQueue k b
mapMaybe = mapMaybeWithKey . const

-- | /O(n)/.  Map values and collect the 'Just' results.
mapMaybeWithKey :: Ord k => (k -> a -> Maybe b) -> MaxPQueue k a -> MaxPQueue k b
mapMaybeWithKey f (MaxPQ q) = MaxPQ (Q.mapMaybeWithKey (f . unDown) q)

-- | /O(n)/.  Map values and separate the 'Left' and 'Right' results.
mapEither :: Ord k => (a -> Either b c) -> MaxPQueue k a -> (MaxPQueue k b, MaxPQueue k c)
mapEither = mapEitherWithKey . const

-- | /O(n)/.  Map values and separate the 'Left' and 'Right' results.
mapEitherWithKey :: Ord k => (k -> a -> Either b c) -> MaxPQueue k a -> (MaxPQueue k b, MaxPQueue k c)
mapEitherWithKey f (MaxPQ q) = case Q.mapEitherWithKey (f . unDown) q of
	(qL, qR) -> (MaxPQ qL, MaxPQ qR)

-- | /O(n)/.  Build a priority queue from the list of (key, value) pairs.
fromList :: Ord k => [(k, a)] -> MaxPQueue k a
fromList = MaxPQ . Q.fromList . fmap (first' Down)

-- | /O(n)/.  Build a priority queue from an ascending list of (key, value) pairs.  /The precondition is not checked./
fromAscList :: [(k, a)] -> MaxPQueue k a
fromAscList = MaxPQ . Q.fromDescList . fmap (first' Down)

-- | /O(n)/.  Build a priority queue from a descending list of (key, value) pairs.  /The precondition is not checked./
fromDescList :: [(k, a)] -> MaxPQueue k a
fromDescList = MaxPQ . Q.fromAscList . fmap (first' Down)

-- | /O(n log n)/.  Return all keys of the queue in ascending order.
keys :: Ord k => MaxPQueue k a -> [k]
keys = fmap fst . toDescList

-- | /O(n log n)/.  Return all elements of the queue in ascending order by key.
elems :: Ord k => MaxPQueue k a -> [a]
elems = fmap snd . toDescList

-- | /O(n log n)/.  Equivalent to 'toDescList'.
assocs :: Ord k => MaxPQueue k a -> [(k, a)]
assocs = toDescList

-- | /O(n log n)/.  Return all (key, value) pairs in ascending order by key.
toAscList :: Ord k => MaxPQueue k a -> [(k, a)]
toAscList (MaxPQ q) = fmap (first' unDown) (Q.toDescList q)

-- | /O(n log n)/.  Return all (key, value) pairs in descending order by key.
toDescList :: Ord k => MaxPQueue k a -> [(k, a)]
toDescList (MaxPQ q) = fmap (first' unDown) (Q.toAscList q)

-- | /O(n log n)/.  Equivalent to 'toAscList'.
-- 
-- If the traversal order is irrelevant, consider using 'toListU'.
toList :: Ord k => MaxPQueue k a -> [(k, a)]
toList = toDescList

-- | /O(n)/.  An unordered right fold over the elements of the queue, in no particular order.
foldrU :: (a -> b -> b) -> b -> MaxPQueue k a -> b
foldrU = foldrWithKeyU . const

-- | /O(n)/.  An unordered right fold over the elements of the queue, in no particular order.
foldrWithKeyU :: (k -> a -> b -> b) -> b -> MaxPQueue k a -> b
foldrWithKeyU f z (MaxPQ q) = Q.foldrWithKeyU (f . unDown) z q

-- | /O(n)/.  An unordered left fold over the elements of the queue, in no particular order.
foldlU :: (b -> a -> b) -> b -> MaxPQueue k a -> b
foldlU f = foldlWithKeyU (const . f)

-- | /O(n)/.  An unordered left fold over the elements of the queue, in no particular order.
foldlWithKeyU :: (b -> k -> a -> b) -> b -> MaxPQueue k a -> b
foldlWithKeyU f z (MaxPQ q) = Q.foldlWithKeyU (\ z -> f z . unDown) z q

-- | /O(n)/.  An unordered traversal over a priority queue, in no particular order.
-- While there is no guarantee in which order the elements are traversed, the resulting
-- priority queue will be perfectly valid.
traverseU :: (Applicative f, Ord b) => (a -> f b) -> MaxPQueue k a -> f (MaxPQueue k b)
traverseU = traverseWithKeyU . const

-- | /O(n)/.  An unordered traversal over a priority queue, in no particular order.
-- While there is no guarantee in which order the elements are traversed, the resulting
-- priority queue will be perfectly valid.
traverseWithKeyU :: (Applicative f, Ord b) => (k -> a -> f b) -> MaxPQueue k a -> f (MaxPQueue k b)
traverseWithKeyU f (MaxPQ q) = MaxPQ <$> Q.traverseWithKeyU (f . unDown) q

-- | /O(n)/.  Return all keys of the queue in no particular order.
keysU :: MaxPQueue k a -> [k]
keysU = fmap fst . toListU

-- | /O(n)/.  Return all elements of the queue in no particular order.
elemsU :: MaxPQueue k a -> [a]
elemsU = fmap snd . toListU

-- | /O(n)/.  Equivalent to 'toListU'.
assocsU :: MaxPQueue k a -> [(k, a)]
assocsU = toListU

-- | /O(n)/.  Returns all (key, value) pairs in the queue in no particular order.
toListU :: MaxPQueue k a -> [(k, a)]
toListU (MaxPQ q) = fmap (first' unDown) (Q.toListU q)

-- | /O(log n)/.  Analogous to @deepseq@ in the @deepseq@ package, but only forces the spine of the binomial heap.
seqSpine :: MaxPQueue k a -> b -> b
seqSpine (MaxPQ q) = Q.seqSpine q