{-# LANGUAGE DeriveGeneric #-}
{-# CFILES cbits/hfrequencyqueue_backend.cpp #-}
{-|
Module      : Data.FrequencyQueue.IO
Description : Provide the IO interface for FrequencyQueue
Copyright   : (c) Andrea Bellandi 2014
License     : GPL-3
Maintainer  : bellaz89@gmai.com
Stability   : experimental
Portability : POSIX

This module export the IO interface of FrequencyQueue. 
-}
module Data.FrequencyQueue.IO(
  -- *Types
  FrequencyQueue(),
  -- *Functions
  -- **Creation functions
  newFrequencyQueue, cloneFrequencyQueue,
  -- **Basic properties
  length, probabilityUnit,
  -- **Pop-push functions
  pushBack, popBack, popBackMax, popBackMin, getRandom, getRandomPop,
  -- **Iterative functions
  mapWprobability, foldWprobability,
  -- **Unsafe interface
  popBackUnsafe, popBackMaxUnsafe, popBackMinUnsafe, getRandomUnsafe, getRandomPopUnsafe) where

import Prelude hiding (length)
import GHC.Generics

import Control.Monad(replicateM)
import Foreign.Concurrent(newForeignPtr)
import Foreign.Marshal.Utils(new)
import Foreign.Marshal.Alloc(free)
import Foreign.CStorable(CStorable, cAlignment, cSizeOf, cPoke, cPeek)
import Foreign.Storable(Storable, alignment, sizeOf, poke, peek)
import Foreign.ForeignPtr(ForeignPtr, withForeignPtr)
import Foreign.StablePtr(StablePtr, deRefStablePtr, freeStablePtr, newStablePtr)
import Foreign.Ptr(Ptr)
import Foreign.C.Types

type FrequencyQueue_ a = Ptr a

-- | FrequencyQueue the basic type of the Library
data FrequencyQueue a = FrequencyQueue{ queue :: ForeignPtr a}

data RandomElement a = RandomElement{ probability :: CUInt,
                                      element :: StablePtr a}
                       deriving(Generic)


instance CStorable (StablePtr a) where
  cAlignment = alignment
  cSizeOf    = sizeOf
  cPoke      = poke
  cPeek      = peek

instance CStorable (RandomElement a)

instance  Storable (RandomElement a) where 
  alignment = cAlignment
  sizeOf    = cSizeOf
  poke      = cPoke
  peek      = cPeek

-- the foreign import shouldn't call functions that call-back the GHC runtime (clone_FrequencyQueue_priv_ and free_FrequencyQueue_priv_)
-- unsafely. Functions that are called unsafely should have constant or constant amortized time to not block the caller OS too much. 
foreign import ccall unsafe new_FrequencyQueue_priv_ :: CUInt -> IO (FrequencyQueue_ a)
foreign import ccall clone_FrequencyQueue_priv_ :: (FrequencyQueue_ a) -> IO (FrequencyQueue_ a)
foreign import ccall unsafe length_priv_ :: (FrequencyQueue_ a) -> IO (CUInt)
foreign import ccall unsafe probability_unit_priv_ :: (FrequencyQueue_ a) -> IO (CUInt)       
foreign import ccall unsafe push_back_priv_ :: (FrequencyQueue_ a) -> (Ptr (RandomElement a)) -> IO ()
foreign import ccall unsafe pop_back_priv_ ::  (FrequencyQueue_ a) -> IO (Ptr (RandomElement a))
foreign import ccall pop_back_max_prob_priv_ :: (FrequencyQueue_ a) -> IO (Ptr (RandomElement a)) 
foreign import ccall pop_back_min_prob_priv_ :: (FrequencyQueue_ a) -> IO (Ptr (RandomElement a)) 
foreign import ccall get_random_priv_ :: (FrequencyQueue_ a) -> IO (Ptr (RandomElement a))
foreign import ccall get_random_pop_priv_ :: (FrequencyQueue_ a) -> IO (Ptr (RandomElement a))
foreign import ccall unsafe reset_iterator_priv_ :: (FrequencyQueue_ a) -> IO ()
foreign import ccall unsafe get_next_priv_ :: (FrequencyQueue_ a) -> IO (Ptr (RandomElement a))
foreign import ccall unsafe get_random_number_priv_ :: (FrequencyQueue_ a) -> IO (CUInt)
foreign import ccall free_FrequencyQueue_priv_ :: FrequencyQueue_ a -> IO ()
foreign import ccall unsafe free_RandomElement_priv_ :: (Ptr (RandomElement a)) -> IO ()
foreign export ccall freeStablePtr :: StablePtr a -> IO () 
foreign export ccall makeNewStableRef :: StablePtr a -> IO (StablePtr a)

makeNewStableRef :: StablePtr a -> IO (StablePtr a)
makeNewStableRef ptr = deRefStablePtr ptr >>= newStablePtr 

-- |Create a new FrequencyQueue with a seed
newFrequencyQueue :: Int -> IO (FrequencyQueue a)
newFrequencyQueue seed = do rawqueue <- new_FrequencyQueue_priv_ (fromIntegral seed)
                            queue_ <- newForeignPtr rawqueue (free_FrequencyQueue_priv_ rawqueue) 
                            return (FrequencyQueue queue_)

-- |Make a clone of the FrequencyQueue Passed
cloneFrequencyQueue :: FrequencyQueue a -> IO (FrequencyQueue a)
cloneFrequencyQueue oldqueue = do rawqueue <- withForeignPtr (queue oldqueue) clone_FrequencyQueue_priv_
                                  queue_ <- newForeignPtr rawqueue (free_FrequencyQueue_priv_ rawqueue) 
                                  return (FrequencyQueue queue_)

-- |Return the number of elements in the queue
length :: FrequencyQueue a -> IO Int
length queue_ =  withForeignPtr (queue queue_) length_priv_ >>= (return . fromIntegral) 

-- |Return the sum of all elements' probabilities passed to the queue
probabilityUnit :: FrequencyQueue a -> IO Int
probabilityUnit queue_ = withForeignPtr (queue queue_) probability_unit_priv_ >>= (return . fromIntegral)

-- |Push an element a in the queue with a corresponding relative probability
pushBack :: FrequencyQueue a -> a -> Int -> IO()
pushBack queue_ element_ probability_ = do stableElement_ <- newStablePtr element_
                                           let cUIntProbability = (fromIntegral probability_)
                                           let randomElement_ = RandomElement cUIntProbability stableElement_
                                           allocatedElement_ <- new randomElement_
                                           withForeignPtr (queue queue_) (\x -> push_back_priv_ x allocatedElement_) 
                                           free allocatedElement_

-- |Pop an element of the queue. Return Nothing if the queue is empty
popBack :: FrequencyQueue a -> IO (Maybe (a,Int))
popBack queue_ = makeSafePop queue_ popBackUnsafe

-- |Pop the element of the queue that have the biggest relative probability.
--  Return Nothing if the queue is empty
popBackMax :: FrequencyQueue a -> IO (Maybe (a,Int))
popBackMax queue_ = makeSafePop queue_ popBackMaxUnsafe

-- |Pop the element of the queue that have the smallest relative probability.
-- Return Nothing if the queue is empty
popBackMin :: FrequencyQueue a -> IO (Maybe (a,Int))
popBackMin queue_ = makeSafePop queue_ popBackMinUnsafe

-- |Return a random element from the queue using its relative probability.
-- Return Nothing if the queue is empty
getRandom :: FrequencyQueue a -> IO (Maybe (a,Int))
getRandom queue_ = makeSafePop queue_ getRandomUnsafe

-- |Pop a random element from the queue using its relative probability.
-- Return Nothing if the queue is empty
getRandomPop :: FrequencyQueue a -> IO (Maybe (a,Int))
getRandomPop queue_ = makeSafePop queue_ getRandomPopUnsafe

-- |Return a new queue with the elements and relative probability mapped
-- by the function provided
mapWprobability :: ((a, Int) -> (b, Int)) -> FrequencyQueue a -> IO (FrequencyQueue b)
mapWprobability fun queue_ = do rnd_number <- withForeignPtr (queue queue_) get_random_number_priv_
                                queue_length <- length queue_
                                newqueue_ <- newFrequencyQueue (fromIntegral rnd_number)
                                withForeignPtr (queue queue_) reset_iterator_priv_
                                replicateM (queue_length) (trasformCopyQueue queue_ newqueue_)
                                return newqueue_
  where
    trasformCopyQueue q1 q2 = do ptr_rawelement <- withForeignPtr (queue q1) get_next_priv_
                                 result <- peek ptr_rawelement
                                 let probability_ = probability result
                                 let elementStable_ = element result
                                 element_ <- deRefStablePtr elementStable_
                                 let transformed_element_ = fun (element_, fromIntegral probability_)
                                 pushBack q2 (fst transformed_element_) (snd transformed_element_)

-- |Return a folded value made by an initial value b and a folding function
-- evaluated on the entire queue.
foldWprobability :: (b -> (a, Int) -> b) -> b -> FrequencyQueue a -> IO b
foldWprobability fold_fun b0 queue_ = do withForeignPtr (queue queue_) reset_iterator_priv_
                                         queue_length <- length queue_
                                         iterateOverFrequencyQueue queue_length b0
  where
    iterateOverFrequencyQueue 0     acc = return acc
    iterateOverFrequencyQueue nitem acc = do ptr_rawelement <- withForeignPtr (queue queue_) get_next_priv_
                                             result <- peek ptr_rawelement
                                             let probability_ = probability result
                                             let elementStable_ = element result
                                             element_ <- deRefStablePtr elementStable_
                                             let next_acc = fold_fun acc (element_, fromIntegral probability_)
                                             iterateOverFrequencyQueue (nitem-1) next_acc

-- |Pop an element of the queue. Fail if empty
popBackUnsafe :: FrequencyQueue a -> IO (a, Int)
popBackUnsafe queue_ = do ptr_rawelement <- withForeignPtr (queue queue_) pop_back_priv_
                          deRefRawElementPtr ptr_rawelement

-- |Pop the element of the queue that have the biggest relative probability.
-- Fail if empty
popBackMaxUnsafe :: FrequencyQueue a -> IO (a, Int)
popBackMaxUnsafe queue_ = do ptr_rawelement <- withForeignPtr (queue queue_) pop_back_max_prob_priv_
                             deRefRawElementPtr ptr_rawelement

-- |Pop the element of the queue that have the smallest relative probability.
-- Fail if empty
popBackMinUnsafe :: FrequencyQueue a -> IO (a, Int)
popBackMinUnsafe queue_ = do ptr_rawelement <- withForeignPtr (queue queue_) pop_back_min_prob_priv_
                             deRefRawElementPtr ptr_rawelement

-- |Pop the element of the queue that have the smallest relative probability.
-- Fail if empty
getRandomUnsafe :: FrequencyQueue a -> IO (a, Int)
getRandomUnsafe queue_ = do ptr_rawelement <- withForeignPtr (queue queue_) get_random_priv_
                            result <- peek ptr_rawelement
                            let probability_ = probability result
                            let elementStable_ = element result
                            element_ <- deRefStablePtr elementStable_
                            return (element_, fromIntegral probability_)

-- |Pop a random element from the queue using its relative probability.
-- Fail if empty
getRandomPopUnsafe :: FrequencyQueue a -> IO (a, Int)
getRandomPopUnsafe queue_ = do ptr_rawelement <- withForeignPtr (queue queue_) get_random_pop_priv_
                               deRefRawElementPtr ptr_rawelement

deRefRawElementPtr :: Ptr (RandomElement a) -> IO (a, Int)
deRefRawElementPtr ptr_rawelement = do result <- peek ptr_rawelement
                                       free_RandomElement_priv_ ptr_rawelement
                                       let probability_ = probability result
                                       let elementStable_ = element result
                                       element_ <- deRefStablePtr elementStable_
                                       freeStablePtr elementStable_
                                       return (element_, fromIntegral probability_)

makeSafePop :: FrequencyQueue a -> (FrequencyQueue a -> IO (a, Int)) -> IO (Maybe (a,Int))
makeSafePop queue_ unsafefun = do qlength <- length queue_ 
                                  if qlength == 0
                                  then return Nothing
                                  else (unsafefun queue_) >>= (\x -> return (Just x))