```{-# LANGUAGE BangPatterns #-}
{-# OPTIONS -Wall #-}
------------------------------------------------------------
-- |
-- Module       : Data.PurePriorityQueue.Internal
--
-- This module exposes the internals of a pure priority queue,
-- implemented on top of "Data.Map".
--
-- Estimates of worst-case time complexity are given.  The value /n/
-- is the number of elements in the queue.  The value /p/ is the
-- cardinality of the set of priorities of the elements in the queue.
-- /p/ is never greater than /n/.
------------------------------------------------------------

module Data.PurePriorityQueue.Internal where

import qualified Data.Map as M
import qualified Data.Foldable as F
import Data.Monoid

import Prelude hiding (filter, null)
import qualified Prelude

-- | A queue of values of type 'a' with priority of type 'p'.
newtype MinMaxQueue p a = MinMaxQueue { unMinMaxQueue :: M.Map p [a] }
deriving (Eq, Ord)

-- | /O(1)/ An empty priority queue.
empty :: MinMaxQueue p a
empty = MinMaxQueue M.empty
{-# INLINE empty #-}

-- | /O(1)/ Test whether a priority queue is empty.
null :: MinMaxQueue p a -> Bool
null = M.null . unMinMaxQueue
{-# INLINE null #-}

-- | /O(1)/ A priority queue with a single entry.
singleton :: Ord p => a -> p -> MinMaxQueue p a
singleton a p = insert a p empty

-- | /O(log p)/ Insert a value with given priority into a priority queue.
insert :: Ord p
=> a
-> p
-> MinMaxQueue p a
-> MinMaxQueue p a
insert a p (MinMaxQueue m) = MinMaxQueue (M.insertWith' (++) p [a] m)
{-# INLINE insert #-}

-- | /O(log p)/ Remove the value with the minimum priority from the
-- queue.
--
-- If the queue is empty, 'deleteMin' returns 'empty'.  Ties are
-- broken arbitrarily.
deleteMin :: Ord p
=> MinMaxQueue p a
-> MinMaxQueue p a
deleteMin m = maybe empty snd (minView m)
{-# INLINE deleteMin #-}

-- | /O(log p)/ Remove the value with the maximum priority from the
-- queue.
--
-- If the queue is empty, 'deleteMax' returns 'empty'.  Ties are
-- broken arbitrarily.
deleteMax :: Ord p
=> MinMaxQueue p a
-> MinMaxQueue p a
deleteMax m = maybe empty snd (maxView m)
{-# INLINE deleteMax #-}

-- | Applies a 'Data.Map.Map' view function to a given priority queue.
viewWith :: (Ord p)
=> (M.Map p [a] -> Maybe ((p, [a]), M.Map p [a]))  -- ^ The view function
-> MinMaxQueue p a                                 -- ^ The priority queue
-> Maybe ((a, p), MinMaxQueue p a)
viewWith f (MinMaxQueue m) = do
((p, a:as), m') <- f m
let m'' = if Prelude.null as
then m'
else M.insert p as m'
return ((a, p), MinMaxQueue m'')
{-# INLINE viewWith #-}

-- | /O(log p)/ View a priority queue to get the (value, priority)
-- pair with the lowest priority and the remainder of the queue.
--
-- Ties are broken arbitrarily.
minView :: Ord p
=> MinMaxQueue p a
-> Maybe ((a, p), MinMaxQueue p a)
minView = viewWith M.minViewWithKey
{-# INLINE minView #-}

-- | /O(log p)/ View a priority queue to get the (value, priority)
-- pair with the highest priority and the remainder of the queue.
--
-- Ties are broken arbitrarily.
maxView :: Ord p => MinMaxQueue p a -> Maybe ((a, p), MinMaxQueue p a)
maxView = viewWith M.maxViewWithKey
{-# INLINE maxView #-}

-- | /O(log p)/ Get the minimum priority of the elements in the queue.
minPriority :: Ord p => MinMaxQueue p a -> Maybe p
minPriority = fmap (snd . fst) . minView
{-# INLINE minPriority #-}

-- | /O(log p)/ Get the maximum priority of the elements in the queue.
maxPriority :: Ord p => MinMaxQueue p a -> Maybe p
maxPriority = fmap (snd . fst) . maxView
{-# INLINE maxPriority #-}

-- | /O(n)/ Fold the priorities and values of a priority queue.
foldWithPriority :: Ord p => (p -> a -> b -> b) -> b -> MinMaxQueue p a -> b
foldWithPriority f s q = M.foldWithKey f' s (unMinMaxQueue q)
where f' !p !vs !acc = foldr (f p) acc vs
{-# INLINE foldWithPriority #-}

-- | /O(log p)/ Split a priority queue 'q' into two queues @(q1, q2)@
-- by the given priority 'p', such that 'q1' contains exactly the
-- entries with priority less than 'p', and 'q2' containes exactly the
-- entries with priority greater than or equal to 'p'.
splitByPriority :: Ord p => p -> MinMaxQueue p a -> (MinMaxQueue p a, MinMaxQueue p a)
splitByPriority p q = (MinMaxQueue lt, MinMaxQueue geq)
where
geq = case meq of
Nothing -> gt
Just eq -> M.insert p eq gt
(lt, meq, gt) = M.splitLookup p (unMinMaxQueue q)
{-# INLINE splitByPriority #-}

-- | /O(n)/ The number of entries in a priority queue.
size :: Ord p => MinMaxQueue p a -> Int
size m = getSum \$ F.foldMap (const \$ Sum 1) m
{-# INLINE size #-}

-- TODO:  filter could be implemented smarter.  Perhaps
-- filterWithPriority also.

-- | /O(n log p)/ Filter all values that satisfy the predicate.
filter :: (Ord p) => (a -> Bool) -> MinMaxQueue p a -> MinMaxQueue p a
filter f = filterWithPriority (\k _ -> f k)
{-# INLINE filter #-}

-- | /O(n log p)/ Filter all entries that satisfy the predicate.
filterWithPriority :: (Ord p)
=> (a -> p -> Bool)
-> MinMaxQueue p a
-> MinMaxQueue p a
filterWithPriority f = foldWithPriority f' empty
where
f' !p !k !q = if f k p then insert k p q else q
{-# INLINE filterWithPriority #-}

-- | /O(n)/ Convert the priority queue into a list of (value,
-- priority) pairs in ascending priority.
--
-- Ties are broken in an arbitrary manner.
toAscList :: (Ord p) => MinMaxQueue p a -> [(a, p)]
toAscList = foldWithPriority (\p a vs -> (a, p) : vs) []
{-# INLINE toAscList #-}

instance Functor (MinMaxQueue p) where
fmap f q = MinMaxQueue (fmap (fmap f) \$ unMinMaxQueue q)
{-# INLINE fmap #-}

instance F.Foldable (MinMaxQueue p) where
foldMap f q = F.foldMap (F.foldMap f) (unMinMaxQueue q)
{-# INLINE foldMap #-}

instance Ord p => Monoid (MinMaxQueue p a) where
mempty = empty
mappend q1 q2 = MinMaxQueue \$ M.unionWith (++) q1' q2'
where
q1' = unMinMaxQueue q1
q2' = unMinMaxQueue q2
{-# INLINE mempty #-}
{-# INLINE mappend #-}
```