{-# LANGUAGE BangPatterns, PatternSynonyms #-}
module Data.MinMaxPSQ
    ( module Data.MinMaxPSQ
    , Binding'
    , pattern Binding
    ) where

import Data.Ord
import qualified Data.Wrapper.PSQ as PSQ
         ;import Data.Wrapper.PSQ as PSQ hiding (insert, insert', null, size)
import Prelude hiding (null, take)

data MinMaxPSQ' k p v = MinMaxPSQ !(PSQ' k p v) !(PSQ' k (Down p) v)
type MinMaxPSQ k p = MinMaxPSQ' k p ()

empty :: MinMaxPSQ' k p v
empty = MinMaxPSQ PSQ.empty PSQ.empty

null :: MinMaxPSQ' k p v -> Bool
null (MinMaxPSQ nq xq) = PSQ.null nq

size :: MinMaxPSQ' k p v -> Int
size (MinMaxPSQ nq xq) = PSQ.size nq

toList :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> [Binding' k p v]
toList (MinMaxPSQ nq xq) = PSQ.toList nq

fromList :: (PSQKey k, Ord p) => [Binding' k p v] -> MinMaxPSQ' k p v
fromList kps = MinMaxPSQ (PSQ.fromList kps)
                         (PSQ.fromList $ map (\(Binding k v p) -> Binding k v (Down p)) kps)

findMin :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> Maybe (Binding' k p v)
findMin (MinMaxPSQ nq xq) = PSQ.findMin nq

findMax :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> Maybe (Binding' k p v)
findMax (MinMaxPSQ nq xq) = fmap (\(Binding k v (Down p)) -> Binding k v p) $ PSQ.findMin xq

insert :: (PSQKey k, Ord p) => k -> p -> MinMaxPSQ k p -> MinMaxPSQ k p
insert k p (MinMaxPSQ nq xq) = MinMaxPSQ (PSQ.insert k p        nq)
                                         (PSQ.insert k (Down p) xq)

insert' :: (PSQKey k, Ord p) => k -> v -> p -> MinMaxPSQ' k p v -> MinMaxPSQ' k p v
insert' k v p (MinMaxPSQ nq xq) = MinMaxPSQ (PSQ.insert' k v p        nq)
                                            (PSQ.insert' k v (Down p) xq)

delete :: (PSQKey k, Ord p) => k -> MinMaxPSQ' k p v -> MinMaxPSQ' k p v
delete k (MinMaxPSQ nq xq) = MinMaxPSQ (PSQ.delete k nq) (PSQ.delete k xq)

deleteMin :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> MinMaxPSQ' k p v
deleteMin (MinMaxPSQ nq xq) = case PSQ.minView nq of
    Just (Binding k _ _, nq') -> MinMaxPSQ nq' (PSQ.delete k xq)
    Nothing                   -> MinMaxPSQ nq xq

deleteMax :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> MinMaxPSQ' k p v
deleteMax (MinMaxPSQ nq xq) = case PSQ.minView xq of
    Just (Binding k _ _, xq') -> MinMaxPSQ (PSQ.delete k nq) xq'
    Nothing                   -> MinMaxPSQ nq xq

minView :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> Maybe (Binding' k p v, MinMaxPSQ' k p v)
minView (MinMaxPSQ nq xq) = fmap (\(Binding k v p, nq') -> (Binding k v p, MinMaxPSQ nq' (PSQ.delete k xq)))
                            $ PSQ.minView nq

maxView :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> Maybe (Binding' k p v, MinMaxPSQ' k p v)
maxView (MinMaxPSQ nq xq) = fmap (\(Binding k v (Down p), xq') -> (Binding k v p, MinMaxPSQ (PSQ.delete k nq) xq'))
                            $ PSQ.minView xq

-- | Maintains a bounded 'MinMaxPSQ' by deleting the maximum element if the
-- insertion would cause the queue to have too many elements.
insertTake :: (PSQKey k, Ord p) => Int -> k -> p -> MinMaxPSQ k p -> MinMaxPSQ k p
insertTake n k p q = take n $ insert k p q

-- | Maintains a bounded 'MinMaxPSQ\'' by deleting the maximum element if the
-- insertion would cause the queue to have too many elements.
insertTake' :: (PSQKey k, Ord p) => Int -> k -> v -> p -> MinMaxPSQ' k p v -> MinMaxPSQ' k p v
insertTake' n k v p q = take n $ insert' k v p q


-- | Truncate the 'MinMaxPSQ' to the given number of lowest priority elements.
take :: (PSQKey k, Ord p) => Int -> MinMaxPSQ' k p v -> MinMaxPSQ' k p v
take !n !q | (size q <= n) = q
           | null q        = q
           | otherwise     = take n $ deleteMax q

-- | Like 'take', except it provides a list deleted bindings.
takeView :: (PSQKey k, Ord p) => Int -> MinMaxPSQ' k p v -> ( [Binding' k p v], MinMaxPSQ' k p v )
takeView !n !q | (size q <= n) = ([], q)
               | null q        = ([], q)
               | otherwise     = let Just (x,q') = maxView q
                                     (xs,q'')    = takeView n q'
                                     ys          = x:xs
                                 in (ys, ys `seq` q'')



lookup' :: (PSQKey k, Ord p) => k -> MinMaxPSQ' k p v -> Maybe (p, v)
lookup' k (MinMaxPSQ q _) = PSQ.lookup k q