{-# LANGUAGE NamedFieldPuns #-}
module Priority (
      Queue(..)
    , empty
    , set
    , toList
  ) where

import Data.Map (Map)
import qualified Data.Map as Map (
    adjust, alter, delete, empty, insert, lookup, lookupMax
  )

data Queue a = Queue {
      byElement :: Map a Int
    , byRank :: Map Int [a]
  }

empty :: Queue a
empty = Queue {byElement = Map.empty, byRank = Map.empty}

set :: Ord a => a -> Int -> Queue a -> Queue a
set x priority queue@(Queue {byElement, byRank}) =
  case Map.lookup x byElement of
    Nothing -> Queue {
          byElement = Map.insert x priority byElement
        , byRank = Map.alter (push x) priority byRank
      }
    Just formerPriority
      | formerPriority >= priority -> queue
      | otherwise -> queue {
            byElement = Map.insert x priority byElement
          , byRank = move x formerPriority priority byRank
        }

push :: a -> Maybe [a] -> Maybe [a]
push a Nothing = Just [a]
push a (Just as) = Just (a:as)

move :: Eq a => a -> Int -> Int -> Map Int [a] -> Map Int [a]
move a from to =
  Map.alter (push a) to . Map.adjust (filter (/= a)) from

toList :: Queue a -> [(Int, a)]
toList = extract . byRank
  where
    extract tmpMap =
      case Map.lookupMax tmpMap of
        Nothing -> []
        Just (k, elems) -> ((,) k <$> elems) ++ extract (Map.delete k tmpMap)