{-# LANGUAGE MultiParamTypeClasses,
             FlexibleInstances,
             GeneralizedNewtypeDeriving,
             NoMonomorphismRestriction #-}

module Data.FingerTree.PSQueue
  (Binding (..), PSQ, size, Data.FingerTree.PSQueue.null, Data.FingerTree.PSQueue.lookup,
   empty, singleton, alter, Data.FingerTree.PSQueue.delete, adjust, adjustWithKey,
   update, updateWithKey, toList, keys, fromList, fromAscList, minView,
   findMin, deleteMin, range, atMost, Data.FingerTree.PSQueue.foldr,
   Data.FingerTree.PSQueue.foldl) where

import qualified Data.Foldable as F
import qualified Data.FingerTree as FT
import Data.FingerTree (FingerTree, ViewL (..), (<|), (|>),
                        Measured (..), split, viewl, (><))
import Data.Monoid
import Data.Ord
import Data.List hiding (lookup)
import Prelude hiding (lookup)

data Prio k a = Prio (Binding k a) | PMax
  deriving (Eq, Ord, Show)

instance (Ord a) => Monoid (Prio k a)
  where mempty = PMax
        mappend PMax y = y
        mappend x PMax = x
        mappend a@(Prio (u :-> x)) b@(Prio (v :-> y)) = if x <= y then a else b

data Key a = NoKey | Key a
  deriving (Eq, Ord, Show)

instance Monoid (Key k)
  where mempty = NoKey
        mappend x NoKey = x
        mappend x y     = y

leqKey k NoKey = False
leqKey k (Key k') = k <= k'

data Binding k p = k :-> p
  deriving (Eq, Ord, Show)

data KPS k p = KPS { kpsKey  :: !(Key k),
                     kpsPrio :: !(Prio k p),
                     kpsSize :: !(Sum Int) }
  deriving (Show)

instance Eq k => Eq (KPS k p) where
  x == y = kpsKey x == kpsKey y

instance Ord k => Ord (KPS k p) where
  compare = comparing kpsKey

instance (Ord p) => Monoid (KPS k p)
  where mempty = KPS mempty mempty mempty
        mappend (KPS k p s) (KPS k' p' s') = KPS (mappend k k')
                                                 (mappend p p')
                                                 (mappend s s')


instance (Ord k, Ord p) => Measured (KPS k p) (Binding k p)
  where measure a@(k :-> p) = KPS (Key k) (Prio a) (Sum 1)

newtype PSQ k p = PSQ (FingerTree (KPS k p) (Binding k p))
  deriving (Eq, Ord, Show, Measured (KPS k p))

-- | O(1). The number of bindings in a queue.
size :: (Ord k, Ord p) => PSQ k p -> Int
size = getSum . kpsSize . measure

-- | O(1). Test if a queue is empty.
null :: (Ord k, Ord p) => PSQ k p -> Bool
null (PSQ q) = FT.null q

-- | O(log n). Determine if a key is in the queue, and its priority.
lookup :: (Ord k, Ord p) => k -> PSQ k p -> Maybe p
lookup k (PSQ q) = let (u,v) = split ((>= Key k) . kpsKey) q
                    in case viewl v of
                         EmptyL -> Nothing
                         (k' :-> p) :< v' | k == k'   -> Just p
                                          | otherwise -> Nothing

-- | O(1). The empty queue.
empty :: (Ord k, Ord p) => PSQ k p
empty = PSQ (FT.empty)

-- | O(1). Construct a queue with a single key/priority binding.
singleton :: (Ord k, Ord p) => k -> p -> PSQ k p
singleton k p = PSQ (FT.singleton (k :-> p))

-- | O(log n). Alters a priority search queue such that @lookup k (alter f k q) = f (lookup k q)@. This can be used to insert, delete,
-- or update a priority in a queue.
alter :: (Ord k, Ord p) => (Maybe p -> Maybe p) -> k -> PSQ k p -> PSQ k p
alter f k (PSQ q) =
  PSQ $ let (u,v) = split (leqKey k . kpsKey) q
         in case viewl v of
              EmptyL -> case f Nothing of
                           Nothing -> q
                           Just p  -> q |> (k :-> p)
              (k' :-> p') :< v'
                | k == k'   -> case f (Just p') of
                                 Nothing -> u >< v'
                                 Just p  -> u >< ((k :-> p) <| v')
                | otherwise -> case f Nothing of
                                 Nothing -> u >< v
                                 Just p  -> u >< ((k :-> p) <| v)

-- | O(log n). Delete a key from a queue.
delete :: (Ord k, Ord p) => k -> PSQ k p -> PSQ k p
delete = alter (const Nothing)

-- | O(log n). Adjust the priority of a key in the queue, provided that key exists.
adjust :: (Ord k, Ord p) => (p -> p) -> k -> PSQ k p -> PSQ k p
adjust f = alter (fmap f)

-- | O(log n). Adjust the priority of a key in the queue, provided that key exists, according to a function which additionally takes
-- the key as a parameter.
adjustWithKey :: (Ord k, Ord p) => (k -> p -> p) -> k -> PSQ k p -> PSQ k p
adjustWithKey f k = adjust (f k) k

-- | O(log n). Update or delete a priority in the queue, provided that key exists.
update :: (Ord k, Ord p) => (p -> Maybe p) -> k -> PSQ k p -> PSQ k p
update f = alter (>>= f) 

-- | O(log n). Update or delete a priority in the queue, provided that key exists, according to a function which additionally takes
-- the key as a parameter.
updateWithKey :: (Ord k, Ord p) => (k -> p -> Maybe p) -> k -> PSQ k p -> PSQ k p
updateWithKey f k = update (f k) k

-- | O(n). Flatten a queue into a list of bindings.
toList :: (Ord k, Ord p) => PSQ k p -> [Binding k p]
toList (PSQ q) = F.toList q

-- | O(n). Extract the list of keys of a queue.
keys :: (Ord k, Ord p) => PSQ k p -> [k]
keys = map (\(k :-> p) -> k) . toList

-- | O(n log n). Construct a queue from a list of bindings.
fromList :: (Ord k, Ord p) => [Binding k p] -> PSQ k p
fromList = PSQ . FT.fromList . sort

-- | O(n log n). Contstruct a queue from an already ascending list of bindings. Does not check that the list is sorted.
fromAscList :: (Ord k, Ord p) => [Binding k p] -> PSQ k p
fromAscList = PSQ . FT.fromList

-- | O(log n). Split a queue into the element with minimum priority, and the remainder.
minView :: (Ord k, Ord p) => PSQ k p -> Maybe (Binding k p, PSQ k p)
minView (PSQ q) =
  let minPrio = kpsPrio . measure $ q
      (u,v) = split ((== minPrio) . kpsPrio) q
   in case viewl v of
        EmptyL -> Nothing
        ((k :-> p) :< v') -> Just (k :-> p, PSQ (u >< v'))

-- | O(1). Find the binding with minimum priority in a queue.
findMin :: (Ord k, Ord p) => PSQ k p -> Maybe (Binding k p)
findMin q = case kpsPrio . measure $ q of
                    PMax   -> Nothing
                    Prio b -> Just b

-- | O(log n). Delete the key with minimum priority from a queue.
deleteMin :: (Ord k, Ord p) => PSQ k p -> PSQ k p
deleteMin q = maybe q id . fmap snd . minView $ q

-- | O(log n). The expression @range (l,u) q@ selects the keys k from q where @l <= k@ and @k <= u@.
range :: (Ord k, Ord p) => (k, k) -> PSQ k p -> PSQ k p
range (l,u) (PSQ q) = PSQ (fst . split ((> Key u) . kpsKey) . snd . split ((>= Key l) . kpsKey) $ q)

-- | O(r (log n)). Finds all the bindings in a queue whose priority is less than the given value.
atMost :: (Ord k, Ord p) => p -> PSQ k p -> [Binding k p]
atMost p (PSQ q) =
  let less (Prio (k :-> p')) = p' < p
      (u,v) = split (less . kpsPrio) q
   in case viewl v of
        EmptyL    -> []
        (b :< v') -> b : atMost p (PSQ v')

-- | Right fold over the list of bindings in a queue. 
foldr :: (Ord k, Ord p) => (Binding k p -> b -> b) -> b -> PSQ k p -> b
foldr f z = Prelude.foldr f z . toList

-- | Left fold over the list of bindings in a queue.
foldl :: (Ord k, Ord p) => (b -> Binding k p -> b) -> b -> PSQ k p -> b
foldl f z = Prelude.foldl f z . toList