{-# LANGUAGE BangPatterns #-}
module Data.MinMaxPSQ where

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

data MinMaxPSQ k p = MinMaxPSQ !(PSQ k p) !(PSQ k (Down p))

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

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

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

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

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

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

findMax :: (PSQKey k, Ord p) => MinMaxPSQ k p -> Maybe (Binding k p)
findMax (MinMaxPSQ nq xq) = fmap (\(k :-> Down p) -> k :-> 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)

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

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

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

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

maxView :: (PSQKey k, Ord p) => MinMaxPSQ k p -> Maybe (Binding k p, MinMaxPSQ k p)
maxView (MinMaxPSQ nq xq) = fmap (\(k :-> Down p, xq') -> (k :-> 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

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