{-# LANGUAGE BangPatterns #-}
--------------------------------------------------------------------
-- |
-- Module    : Data.Vector.Random.Mersenne
-- Copyright : (c) Don Stewart 2010

-- License   : BSD3
--
-- Maintainer: Don Stewart <dons@galois.com>
-- Stability : provisional
-- Portability:
--
-- Generator for vectors of differen types
--

module Data.Vector.Random.Mersenne (

    -- * Fill a vector with randoms.
    randoms,

    -- * Class of MT random types
    PureMTRandom(..)

  ) where

-- import qualified System.Random.Mersenne     as R
import qualified System.Random.Mersenne.Pure64 as R


import qualified Data.Vector.Generic         as G
import qualified Data.Vector.Generic.Mutable as GM

import qualified Data.Vector.Fusion.Stream         as Stream
import qualified Data.Vector.Fusion.Stream.Monadic as S
import qualified Data.Vector.Fusion.Stream.Size    as S
import Data.Vector.Fusion.Util

-- import System.IO.Unsafe
-- import GHC.IOBase
import Data.Bits
import Data.Int
import Data.Word

-- | Return a random vector of length @n@, filled with random elements
-- of type @a@ generated by the mersenne-twister.
--
-- E.g. to compute the sum of 100 million random Double values in a vector:
--
-- > import qualified Data.VectorUnboxed as U
-- > import System.Random.Mersenne.Pure64
-- > import qualified Data.Vector.Random.Mersenne as G
-- > 
-- > main = do
-- >     g <- newPureMT
-- >     let a = G.random g 10000000 :: U.Vector Double
-- >     print (U.sum a)
--
-- The generator will fuse under stream fusion, so e.g. sum . random g
-- will allocate no intermediate array.
--

randoms :: (PureMTRandom a, G.Vector v a) => R.PureMT -> Int -> v a
randoms g n = G.unstream (randomS g n)
{-# INLINE randoms #-}

-- A stream of 'n' random numbers.
--
randomS :: PureMTRandom a => R.PureMT -> Int -> Stream.Stream a
{-# INLINE [1] randomS #-}
randomS g n = S.Stream (return . step) (n, g) (S.Exact (delay_inline max n 0))
  where
    {-# INLINE [0] step #-}
    step (i,g) | i <= 0    = S.Done
               | otherwise = g `seq` case random g of
                                (r, g') -> S.Yield r (i-1, g')

------------------------------------------------------------------------

-- | Class of types that we have efficient generators for.
class PureMTRandom a where

    -- | Given a pure mersenne twister state, yield a new random value,
    -- and the next state.
    random :: R.PureMT -> (a, R.PureMT)
    {-# INLINE random #-}

instance PureMTRandom Double where
    random = R.randomDouble
    {-# INLINE random #-}

-- XXX this doesn't inline properly under ghc 6.12
instance PureMTRandom Float  where
    random !g = case R.randomDouble g of
                    (i, g') -> (realToFrac i, g')
    {-# INLINE random #-}

instance PureMTRandom Word where
    random = R.randomWord
    {-# INLINE random #-}

instance PureMTRandom Word8 where
    random g = case R.randomWord g of (i, g') -> (fromIntegral i, g')
    {-# INLINE random #-}

instance PureMTRandom Word16 where
    random g = case R.randomWord g of (i, g') -> (fromIntegral i, g')
    {-# INLINE random #-}

instance PureMTRandom Word32 where
    random g = case R.randomWord g of (i, g') -> (fromIntegral i, g')
    {-# INLINE random #-}

instance PureMTRandom Word64 where
    random = R.randomWord64
    {-# INLINE random #-}

instance PureMTRandom Int where
    random = R.randomInt
    {-# INLINE random #-}

instance PureMTRandom Int8 where
    random g = case R.randomInt g of (i, g') -> (fromIntegral i, g')
    {-# INLINE random #-}

instance PureMTRandom Int16 where
    random g = case R.randomInt g of (i, g') -> (fromIntegral i, g')
    {-# INLINE random #-}

instance PureMTRandom Int32 where
    random g = case R.randomInt g of (i, g') -> (fromIntegral i, g')
    {-# INLINE random #-}

instance PureMTRandom Int64 where
    random = R.randomInt64
    {-# INLINE random #-}

instance PureMTRandom Bool where
    random g = case R.randomWord g of (i, g') -> (i .&. i /= 0, g')
    {-# INLINE random #-}

instance PureMTRandom Integer where
    random g = case R.randomInt g of (i, g') -> (fromIntegral i, g')
    {-# INLINE random #-}

------------------------------------------------------------------------

{-
random g n = do
    v  <- GM.new n
    fill v 0
    G.unsafeFreeze v
  where
    fill v i
        | i < n = do
            x <- R.random g
            GM.unsafeWrite v i x
            fill v (i+1)
        | otherwise = return ()
-}

------------------------------------------------------------------------

{-
-- A stream of random 'n' random numbers.
--
randomS :: (R.MTRandom a) => R.PureMT -> Int -> Stream.Stream a
{-# INLINE [1] randomS #-}
randomS g n = S.Stream (return . step) n (S.Exact (delay_inline max n 0))
  where
    {-# INLINE [0] step #-}
    step i | i <= 0    = S.Done
           | otherwise =
               case unsafeDupablePerformIO (R.random' i g) of
                    r -> S.Yield r (i-1)


random' :: R.MTRandom a => Int -> R.PureMT -> IO a
{-# NOINLINE random' #-}
random' _ gen = R.random gen

-}