module Data.InfinitePriorityQueue where

import Control.Applicative
import Data.Monoid

import Data.Tuple.HT(mapFst, mapSnd)

newtype InfPQ a = InfPQ { unInfPQ :: Int -> [(Int, a)] } 

inf = 10000000

instance (Show a) => Show (InfPQ a) where
  show (InfPQ a) = show $ a inf

-- alternate :: [a] -> [a] -> [a]
-- alternate [] l = l
-- alternate l [] = l
-- alternate (a : as) (b:bs) = a:b:(alternate as bs)

instance Functor InfPQ where
  fmap f (InfPQ as) = InfPQ $ fmap (mapSnd f) . as

takeSorted :: [[(Int,a)]] -> [(Int,a)]
takeSorted [] = []
takeSorted ([]:rs) = takeSorted rs
takeSorted (r:[]) = r
takeSorted (r:rs) =  
  let srs = takeSorted rs
      (a:as) = r
  in case as of [] -> a : takeSorted rs
                _ -> let lim = fst $ head as
                         srslt = takeWhile ((<lim).fst) srs
                         srsgt = dropWhile ((<lim).fst) srs
                     in (a : srslt) ++ merge as srsgt

merge :: [(Int,a)] -> [(Int,a)] -> [(Int,a)]
merge [] rs = rs
merge (l:ls) [] = l:ls
merge ((pl,vl):ls) ((pr,vr):rs) | pl <= pr = (pl,vl): merge ls ((pr,vr):rs)
merge ((pl,vl):ls) ((pr,vr):rs) | otherwise = (pr,vr): merge ((pl,vl):ls) rs

mapFstSnd :: (a -> c) -> (b->d) -> (a,b) -> (c,d)
mapFstSnd fa fb (a,b) = (fa a, fb b)

instance Applicative InfPQ where
  pure = singleton 0
  (InfPQ as) <*> (InfPQ bs) = InfPQ $ \lim -> takeSorted $
                                           do (p,v) <- as lim
                                              return $ map (mapFstSnd (+p) v) $ bs $ lim - p

instance Monad InfPQ where
  return = pure
  (InfPQ as) >>= f = InfPQ $ \lim -> takeSorted $ map (\(p,v) -> map (mapFst (+p)) $ unInfPQ (f v) (lim - p)) $ as lim
  fail _ = empty

instance Alternative InfPQ where
  empty = InfPQ $ const []
  (InfPQ a) <|> (InfPQ b) = InfPQ $ \c -> merge (a c) (b c)

fromAscList :: [(Int, a)] -> InfPQ a
fromAscList = InfPQ . const
toAscList :: Int -> InfPQ a -> [(Int, a)]
toAscList = flip unInfPQ

singleton :: Int -> a -> InfPQ a
singleton p v = fromAscList [(p,v)]

increasePriorities :: Int -> InfPQ v -> InfPQ v
increasePriorities dp (InfPQ l) = InfPQ $ \lim -> if lim < dp then [] else map (mapFst (+dp)) $ l (lim - dp)

sizeUnderPriority :: Int -> InfPQ v -> Int
sizeUnderPriority p q = length $ toAscList p q