-- Priority queues, implemented as a partially ordered heap.
-- An element can occur more than once in a queue.
-- The "top" element is the one that sorts least.
module Util.PQueue (
  PQueue,
  Priority(..),  -- comparison class

  -- construction
  emptyPQ, unitPQ,   -- O(1)
  listToPQ,          -- O(length list)

  -- use
  nextPQ, popPQ,
  peekPQ,            -- O(1)
  
  -- adding elements (fast)
  addToPQ,           -- O(log(size PQ))
  addListToPQ,       -- O(log(size PQ + length list) * (length list))

  -- deleting elements (slow)
  -- "delete" deletes only one occurrence, "purge" deletes all occurrences.
  delFromPQ,
--  purgeFromPQ,          -- O(size PQ)
--  delListFromPQ,
--  purgeListFromPQ,      -- O(size PQ * length list)

  -- combination
--  plusPQ,    -- O(log(size biggerPQ) * (size smallerPQ))
--  minusPQ, intersectPQ,  -- O(size PQ1 * size PQ2)

  -- conversion
  orderedListPQ, listPQ, filterPQ, foldPQ,  -- O(size PQ)
--  mapPQ,     -- FIXME: cost unknown

  -- interrogation
  sizePQ, isEmptyPQ,   -- O(1)

  -- lookup (slow)
--  elemPQ,  -- O(size PQ)
) where

-- We use a comparison that's separate from Ord, because Ord implies Eq,
-- and we want to make a distinction between element priority and element
-- identity, so that two different elements can have the same priority.
-- (The distinction does not matter for PQueue itself, but it's useful
-- for callers.)
class Priority p where
  pcompare :: p -> p -> Ordering

-- PQueue invariants:
-- * A key in a BranchPQ is higher priority than all the keys in its subtrees
-- * The size field in BranchPQ is equal to the number of elements in its tree.
data PQueue elt = BranchPQ elt Int (PQueue elt) (PQueue elt)
                | EmptyPQ
-- I tried adding a UnitPQ constructor but that just slowed things down.
-- (It did reduce allocation, though)

emptyPQ :: PQueue elt
emptyPQ = EmptyPQ

-- Create a PQueue with one element
unitPQ :: elt -> PQueue elt
unitPQ elt = BranchPQ elt 1 EmptyPQ EmptyPQ

-- Convert a list to a PQueue
listToPQ :: (Priority elt) => [elt] -> PQueue elt
listToPQ = addListToPQ emptyPQ

-- Extract the highest priority element from the PQueue
nextPQ :: (Priority elt) => PQueue elt -> (elt, PQueue elt)
nextPQ pq = (peekPQ pq, popPQ pq)

-- Return the highest priority element of the PQueue
peekPQ :: PQueue elt -> elt
peekPQ EmptyPQ = error "peekPQ: empty PQueue"
peekPQ (BranchPQ elt _ _ _) = elt

-- Discard the highest priority element of the PQueue, if any
popPQ :: (Priority elt) => PQueue elt -> PQueue elt
popPQ EmptyPQ = EmptyPQ
popPQ (BranchPQ _ _ EmptyPQ right) = right
popPQ (BranchPQ _ _ left EmptyPQ) = left
popPQ (BranchPQ _ size left right) =
  case peekPQ left `pcompare` peekPQ right of
    GT -> descend_left
    LT -> descend_right
    EQ -> case sizePQ left `compare` sizePQ right of
      LT -> descend_right  -- make smallest tree smaller, but be done faster
      _ -> descend_left    -- idem
  where
    descend_left = BranchPQ (peekPQ left) (size-1) (popPQ left) right
    descend_right = BranchPQ (peekPQ right) (size-1) left (popPQ right)

addToPQ :: (Priority elt) => PQueue elt -> elt -> PQueue elt
addToPQ EmptyPQ elt = unitPQ elt
addToPQ (BranchPQ top size left right) elt = 
  case elt `pcompare` top of
    GT -> addToPQ (BranchPQ elt size left right) top
    _ -> case sizePQ left `compare` sizePQ right of
      LT -> BranchPQ top (size+1) (addToPQ left elt) right
      _ -> BranchPQ top (size+1) left (addToPQ right elt)

addListToPQ :: (Priority elt) => PQueue elt -> [elt] -> PQueue elt
addListToPQ = foldl addToPQ

delFromPQ :: (Eq elt, Priority elt) => PQueue elt -> elt -> PQueue elt
delFromPQ EmptyPQ _ = EmptyPQ
delFromPQ pq@(BranchPQ top size left right) old =
  -- Check if we found it.
  if top == old
    then popPQ pq
    -- Check if old can be a child of top
    else case top `pcompare` old of
      LT -> pq
      _ -> let newleft = delFromPQ left old
               newright = delFromPQ right old
           -- Check if it's a left-hand child
           in if sizePQ newleft /= sizePQ left
                then BranchPQ top (size-1) newleft right
                -- Check if it's a right-hand child
                else if sizePQ newright /= sizePQ right
                       then BranchPQ top (size - 1) left newright
                       else pq

{-# INLINE sizePQ #-}
sizePQ :: PQueue elt -> Int
sizePQ EmptyPQ = 0
sizePQ (BranchPQ _ size _ _) = size

isEmptyPQ :: PQueue elt -> Bool
isEmptyPQ EmptyPQ = True
isEmptyPQ (BranchPQ _ _ _ _) = False

listPQ :: PQueue elt -> [elt]
listPQ EmptyPQ = []
listPQ (BranchPQ top _ left right) = (top : listPQ left ++ listPQ right)

orderedListPQ :: (Priority elt) => PQueue elt -> [elt]
orderedListPQ EmptyPQ = []
orderedListPQ (BranchPQ top _ left right) =
  top : merge (orderedListPQ left) (orderedListPQ right)
  where merge [] ys = ys
        merge xs [] = xs
        merge (x : xs) (y : ys) = case x `pcompare` y of
          GT -> x : merge xs (y : ys)
          LT -> y : merge (x : xs) ys
          EQ -> x : y : merge xs ys

filterPQ :: (Priority elt) => (elt -> Bool) -> PQueue elt -> PQueue elt
filterPQ _ EmptyPQ = EmptyPQ
filterPQ f pq@(BranchPQ top _ left right) =
  if (f top)
    then BranchPQ top (sizePQ left' + sizePQ right' + 1) left' right'
    else filterPQ f (popPQ pq)
  where left' = filterPQ f left
        right' = filterPQ f right

foldPQ :: (elt -> a -> a) -> a -> PQueue elt -> a
foldPQ _ z EmptyPQ = z
foldPQ f z (BranchPQ top _ left right) =
  foldPQ f (f top (foldPQ f z left)) right