{-# LANGUAGE FlexibleContexts #-}

-- |
-- Module     : Simulation.Aivika.PriorityQueue
-- Copyright  : Copyright (c) 2009-2017, David Sorokin <david.sorokin@gmail.com>
-- License    : BSD3
-- Maintainer : David Sorokin <david.sorokin@gmail.com>
-- Stability  : experimental
-- Tested with: GHC 8.0.1
--
-- An imperative heap-based priority queue to implement the event queue.
--
module Simulation.Aivika.PriorityQueue.EventQueue
       (PriorityQueue, 
        Priority,
        queueNull, 
        queueCount,
        newQueue, 
        enqueue, 
        dequeue, 
        queueFront) where 

import Data.Array
import Data.Array.MArray.Safe
import Data.Array.IO.Safe
import Data.IORef
import Data.Maybe

import Control.Monad

-- | The priority value (greater is higher).
type Priority = Int

-- | The 'PriorityQueue' type represents an imperative heap-based 
-- priority queue.
data PriorityQueue a = 
  PriorityQueue { pqKeys  :: IORef (IOUArray Int Double),
                  pqPris  :: IORef (IOUArray Int Priority),
                  pqVals  :: IORef (IOArray Int a),
                  pqSize  :: IORef Int }

increase :: PriorityQueue a -> Int -> IO ()
increase pq capacity = 
  do let keyRef = pqKeys pq
         priRef = pqPris pq
         valRef = pqVals pq
     keys <- readIORef keyRef
     pris <- readIORef priRef
     vals <- readIORef valRef
     (il, iu)  <- getBounds keys
     let len = (iu - il) + 1
         capacity' | len < 64  = max capacity ((len + 1) * 2)
                   | otherwise = max capacity ((len `div` 2) * 3)
         il' = il
         iu' = il + capacity' - 1
     keys' <- newArray_ (il', iu')
     pris' <- newArray_ (il', iu')
     vals' <- newArray_ (il', iu')
     mapM_ (\i -> do { k <- readArray keys i; writeArray keys' i k }) [il..iu]
     mapM_ (\i -> do { p <- readArray pris i; writeArray pris' i p }) [il..iu]
     mapM_ (\i -> do { v <- readArray vals i; writeArray vals' i v }) [il..iu]
     writeIORef keyRef keys'
     writeIORef priRef pris'
     writeIORef valRef vals'

siftUp :: IOUArray Int Double
       -> IOUArray Int Priority
       -> IOArray Int a
       -> Int
       -> Double
       -> Priority
       -> a 
       -> IO ()
siftUp keys pris vals i k p v = loop i
  where loop i =
          if i == 0 
          then do writeArray keys i k
                  writeArray pris i p
                  writeArray vals i v
          else do let n = (i - 1) `div` 2
                  kn <- readArray keys n
                  pn <- readArray pris n
                  if gte k p kn pn    -- (k, -p) >= (kn, -pn)
                    then do writeArray keys i k
                            writeArray pris i p
                            writeArray vals i v
                    else do pn <- readArray pris n
                            vn <- readArray vals n
                            writeArray keys i kn
                            writeArray pris i pn
                            writeArray vals i vn
                            loop n

siftDown :: IOUArray Int Double
         -> IOUArray Int Priority
         -> IOArray Int a
         -> Int
         -> Int
         -> Double
         -> Priority
         -> a 
         -> IO ()
siftDown keys pris vals size i k p v = loop i
  where loop i =
          if i >= (size `div` 2)
          then do writeArray keys i k
                  writeArray pris i p
                  writeArray vals i v
          else do let n  = 2 * i + 1
                      n' = n + 1
                  kn  <- readArray keys n
                  pn  <- readArray pris n
                  if n' >= size 
                    then if lte k p kn pn    -- (k, -p) <= (kn, -pn)
                    then do writeArray keys i k
                            writeArray pris i p
                            writeArray vals i v
                    else do pn <- readArray pris n
                            vn <- readArray vals n
                            writeArray keys i kn
                            writeArray pris i pn
                            writeArray vals i vn
                            loop n
                    else do kn' <- readArray keys n'
                            pn' <- readArray pris n'
                            -- (kn, -pn) > (kn', -pn')
                            let n''  = if gt kn pn kn' pn' then n' else n
                                kn'' = if n'' == n' then kn' else kn
                                pn'' = if n'' == n' then pn' else pn
                            if lte k p kn'' pn''    -- (k, -p) <= (kn'', -pn'')
                              then do writeArray keys i k
                                      writeArray pris i p
                                      writeArray vals i v
                              else do pn'' <- readArray pris n''
                                      vn'' <- readArray vals n''
                                      writeArray keys i kn''
                                      writeArray pris i pn''
                                      writeArray vals i vn''
                                      loop n''

-- | Test whether the priority queue is empty.
queueNull :: PriorityQueue a -> IO Bool
queueNull pq =
  do size <- readIORef (pqSize pq)
     return $ size == 0

-- | Return the number of elements in the priority queue.
queueCount :: PriorityQueue a -> IO Int
queueCount pq = readIORef (pqSize pq)

-- | Create a new priority queue.
newQueue :: IO (PriorityQueue a)
newQueue =
  do keys <- newArray_ (0, 10)
     pris <- newArray_ (0, 10)
     vals <- newArray_ (0, 10)
     keyRef  <- newIORef keys
     priRef  <- newIORef pris
     valRef  <- newIORef vals
     sizeRef <- newIORef 0
     return PriorityQueue { pqKeys = keyRef,
                            pqPris = priRef,
                            pqVals = valRef, 
                            pqSize = sizeRef }

-- | Enqueue a new element with the specified priority.
enqueue :: PriorityQueue a -> Double -> Priority -> a -> IO ()
enqueue pq k p v =
  do i <- readIORef (pqSize pq)
     keys <- readIORef (pqKeys pq)
     (il, iu) <- getBounds keys
     when (i >= iu - il) $ increase pq (i + 2)  -- plus one element on the end
     writeIORef (pqSize pq) (i + 1)
     keys <- readIORef (pqKeys pq)  -- it can be another! (side-effect)
     pris <- readIORef (pqPris pq)
     vals <- readIORef (pqVals pq)
     siftUp keys pris vals i k p v

-- | Dequeue the element with the minimal priority.
dequeue :: PriorityQueue a -> IO ()
dequeue pq =
  do size <- readIORef (pqSize pq)
     when (size == 0) $ error "Empty priority queue: dequeue"
     let i = size - 1
     writeIORef (pqSize pq) i
     keys <- readIORef (pqKeys pq)
     pris <- readIORef (pqPris pq)
     vals <- readIORef (pqVals pq)
     k  <- readArray keys i
     p  <- readArray pris i
     v  <- readArray vals i
     let k0 = 0.0
         p0 = 0
         v0 = undefined
     -- k0 <- readArray keys size
     -- p0 <- readArray pris size
     -- v0 <- readArray vals size
     writeArray keys i k0
     writeArray pris i p0
     writeArray vals i v0
     when (i > 0) $
       siftDown keys pris vals i 0 k p v

-- | Return the element with the minimal priority.
queueFront :: PriorityQueue a -> IO (Double, Priority, a)
queueFront pq =
  do size <- readIORef (pqSize pq)
     when (size == 0) $ error "Empty priority queue: front"
     keys <- readIORef (pqKeys pq)
     pris <- readIORef (pqPris pq)
     vals <- readIORef (pqVals pq)
     k <- readArray keys 0
     p <- readArray pris 0
     v <- readArray vals 0
     return (k, p, v)

-- | Whether the first pair is greater than the second one.
gt :: Double -> Priority -> Double -> Priority -> Bool
{-# INLINE gt #-}
gt k1 p1 k2 p2 = (k1 > k2) || (k1 == k2 && p1 < p2)

-- | Whether the first pair is greater than or equal to the second one.
gte :: Double -> Priority -> Double -> Priority -> Bool
{-# INLINE gte #-}
gte k1 p1 k2 p2 = (k1 > k2) || (k1 == k2 && p1 <= p2)

-- | Whether the first pair is less than the second one.
lt :: Double -> Priority -> Double -> Priority -> Bool
{-# INLINE lt #-}
lt k1 p1 k2 p2 = gt k2 p2 k1 p1

-- | Whether the first pair is less than or equal to the second one.
lte :: Double -> Priority -> Double -> Priority -> Bool
{-# INLINE lte #-}
lte k1 p1 k2 p2 = gte k2 p2 k1 p1