{-# LANGUAGE TypeFamilies #-}

-- |
-- Module     : Simulation.Aivika.IO.Generator
-- 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
--
-- Here is defined a random number generator, where
-- the 'IO' monad is an instance of 'MonadGenerator'.
--
module Simulation.Aivika.IO.Generator () where

import System.Random
import qualified System.Random.MWC as MWC

import Control.Monad
import Control.Monad.Trans

import Data.IORef
import Data.Vector
import Data.Functor

import Simulation.Aivika.Trans.Generator
import Simulation.Aivika.Trans.Generator.Primitive

instance MonadGenerator IO where
-- instance (Functor m, Monad m, MonadIO m, MonadTemplate m) => MonadGenerator m where

  {-# SPECIALISE instance MonadGenerator IO #-}

  data Generator IO =
    Generator { Generator IO -> IO Double
generator01 :: IO Double,
                -- ^ the generator of uniform numbers from 0 to 1
                Generator IO -> IO Double
generatorNormal01 :: IO Double,
                -- ^ the generator of normal numbers with mean 0 and variance 1
                Generator IO -> IO Int
generatorSequenceNo :: IO Int
                -- ^ the generator of sequence numbers
              }

  {-# INLINE generateUniform #-}
  generateUniform :: Generator IO -> Double -> Double -> IO Double
generateUniform = forall (m :: * -> *).
Monad m =>
m Double -> Double -> Double -> m Double
generateUniform01 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Generator IO -> IO Double
generator01

  {-# INLINE generateUniformInt #-}
  generateUniformInt :: Generator IO -> Int -> Int -> IO Int
generateUniformInt = forall (m :: * -> *). Monad m => m Double -> Int -> Int -> m Int
generateUniformInt01 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Generator IO -> IO Double
generator01

  {-# INLINE generateTriangular #-}
  generateTriangular :: Generator IO -> Double -> Double -> Double -> IO Double
generateTriangular = forall (m :: * -> *).
Monad m =>
m Double -> Double -> Double -> Double -> m Double
generateTriangular01 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Generator IO -> IO Double
generator01

  {-# INLINE generateNormal #-}
  generateNormal :: Generator IO -> Double -> Double -> IO Double
generateNormal = forall (m :: * -> *).
Monad m =>
m Double -> Double -> Double -> m Double
generateNormal01 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Generator IO -> IO Double
generatorNormal01

  {-# INLINE generateLogNormal #-}
  generateLogNormal :: Generator IO -> Double -> Double -> IO Double
generateLogNormal = forall (m :: * -> *).
Monad m =>
m Double -> Double -> Double -> m Double
generateLogNormal01 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Generator IO -> IO Double
generatorNormal01

  {-# INLINE generateExponential #-}
  generateExponential :: Generator IO -> Double -> IO Double
generateExponential = forall (m :: * -> *). Monad m => m Double -> Double -> m Double
generateExponential01 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Generator IO -> IO Double
generator01

  {-# INLINE generateErlang #-}
  generateErlang :: Generator IO -> Double -> Int -> IO Double
generateErlang = forall (m :: * -> *).
Monad m =>
m Double -> Double -> Int -> m Double
generateErlang01 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Generator IO -> IO Double
generator01

  {-# INLINE generatePoisson #-}
  generatePoisson :: Generator IO -> Double -> IO Int
generatePoisson = forall (m :: * -> *). Monad m => m Double -> Double -> m Int
generatePoisson01 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Generator IO -> IO Double
generator01

  {-# INLINE generateBinomial #-}
  generateBinomial :: Generator IO -> Double -> Int -> IO Int
generateBinomial = forall (m :: * -> *). Monad m => m Double -> Double -> Int -> m Int
generateBinomial01 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Generator IO -> IO Double
generator01

  {-# INLINE generateGamma #-}
  generateGamma :: Generator IO -> Double -> Double -> IO Double
generateGamma Generator IO
g = forall (m :: * -> *).
Monad m =>
m Double -> m Double -> Double -> Double -> m Double
generateGamma01 (Generator IO -> IO Double
generatorNormal01 Generator IO
g) (Generator IO -> IO Double
generator01 Generator IO
g)

  {-# INLINE generateBeta #-}
  generateBeta :: Generator IO -> Double -> Double -> IO Double
generateBeta Generator IO
g = forall (m :: * -> *).
Monad m =>
m Double -> m Double -> Double -> Double -> m Double
generateBeta01 (Generator IO -> IO Double
generatorNormal01 Generator IO
g) (Generator IO -> IO Double
generator01 Generator IO
g)

  {-# INLINE generateWeibull #-}
  generateWeibull :: Generator IO -> Double -> Double -> IO Double
generateWeibull = forall (m :: * -> *).
Monad m =>
m Double -> Double -> Double -> m Double
generateWeibull01 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Generator IO -> IO Double
generator01

  {-# INLINE generateDiscrete #-}
  generateDiscrete :: forall a. Generator IO -> DiscretePDF a -> IO a
generateDiscrete = forall (m :: * -> *) a. Monad m => m Double -> DiscretePDF a -> m a
generateDiscrete01 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Generator IO -> IO Double
generator01

  {-# INLINE generateSequenceNo #-}
  generateSequenceNo :: Generator IO -> IO Int
generateSequenceNo = Generator IO -> IO Int
generatorSequenceNo 

  {-# INLINABLE newGenerator #-}
  newGenerator :: GeneratorType IO -> IO (Generator IO)
newGenerator GeneratorType IO
tp =
    case GeneratorType IO
tp of
      GeneratorType IO
SimpleGenerator ->
        forall a (m :: * -> *).
(Variate a, PrimMonad m) =>
Gen (PrimState m) -> m a
MWC.uniform forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO GenIO
MWC.createSystemRandom forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (m :: * -> *).
MonadGenerator m =>
m Double -> m (Generator m)
newRandomGenerator01
      SimpleGeneratorWithSeed Word32
x ->
        forall a (m :: * -> *).
(Variate a, PrimMonad m) =>
Gen (PrimState m) -> m a
MWC.uniform forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) (v :: * -> *).
(PrimMonad m, Vector v Word32) =>
v Word32 -> m (Gen (PrimState m))
MWC.initialize (forall a. a -> Vector a
singleton Word32
x) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (m :: * -> *).
MonadGenerator m =>
m Double -> m (Generator m)
newRandomGenerator01
      CustomGenerator IO (Generator IO)
g ->
        IO (Generator IO)
g
      CustomGenerator01 IO Double
g ->
        forall (m :: * -> *).
MonadGenerator m =>
m Double -> m (Generator m)
newRandomGenerator01 IO Double
g

  {-# INLINABLE newRandomGenerator #-}
  newRandomGenerator :: forall g. RandomGen g => g -> IO (Generator IO)
newRandomGenerator g
g = 
    do IORef g
r <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. a -> IO (IORef a)
newIORef g
g
       let g01 :: IO Double
g01 = do g
g <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO a
readIORef IORef g
r
                    let (Double
x, g
g') = forall a g. (Random a, RandomGen g) => g -> (a, g)
random g
g
                    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef IORef g
r g
g'
                    forall (m :: * -> *) a. Monad m => a -> m a
return Double
x
       forall (m :: * -> *).
MonadGenerator m =>
m Double -> m (Generator m)
newRandomGenerator01 IO Double
g01

  {-# INLINABLE newRandomGenerator01 #-}
  newRandomGenerator01 :: IO Double -> IO (Generator IO)
newRandomGenerator01 IO Double
g01 =
    do IO Double
gNormal01 <- forall (m :: * -> *). MonadIO m => m Double -> m (m Double)
newNormalGenerator01 IO Double
g01
       IORef Int
gSeqNoRef <- forall a. a -> IO (IORef a)
newIORef Int
0
       let gSeqNo :: IO Int
gSeqNo = do { Int
x <- forall a. IORef a -> IO a
readIORef IORef Int
gSeqNoRef; forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef Int
gSeqNoRef (forall a. Num a => a -> a -> a
+Int
1); forall (m :: * -> *) a. Monad m => a -> m a
return Int
x }
       forall (m :: * -> *) a. Monad m => a -> m a
return Generator { generator01 :: IO Double
generator01 = IO Double
g01,
                          generatorNormal01 :: IO Double
generatorNormal01 = IO Double
gNormal01,
                          generatorSequenceNo :: IO Int
generatorSequenceNo = IO Int
gSeqNo }

-- | Create a normal random number generator with mean 0 and variance 1
-- by the specified generator of uniform random numbers from 0 to 1.
newNormalGenerator01 :: MonadIO m
                        => m Double
                        -- ^ the generator
                        -> m (m Double)
{-# INLINABLE newNormalGenerator01 #-}
newNormalGenerator01 :: forall (m :: * -> *). MonadIO m => m Double -> m (m Double)
newNormalGenerator01 m Double
g =
  do IORef Double
nextRef <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. a -> IO (IORef a)
newIORef Double
0.0
     IORef Bool
flagRef <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. a -> IO (IORef a)
newIORef Bool
False
     IORef Double
xi1Ref  <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. a -> IO (IORef a)
newIORef Double
0.0
     IORef Double
xi2Ref  <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. a -> IO (IORef a)
newIORef Double
0.0
     IORef Double
psiRef  <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. a -> IO (IORef a)
newIORef Double
0.0
     let loop :: m ()
loop =
           do Double
psi <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO a
readIORef IORef Double
psiRef
              if (Double
psi forall a. Ord a => a -> a -> Bool
>= Double
1.0) Bool -> Bool -> Bool
|| (Double
psi forall a. Eq a => a -> a -> Bool
== Double
0.0)
                then do Double
g1 <- m Double
g
                        Double
g2 <- m Double
g
                        let xi1 :: Double
xi1 = Double
2.0 forall a. Num a => a -> a -> a
* Double
g1 forall a. Num a => a -> a -> a
- Double
1.0
                            xi2 :: Double
xi2 = Double
2.0 forall a. Num a => a -> a -> a
* Double
g2 forall a. Num a => a -> a -> a
- Double
1.0
                            psi :: Double
psi = Double
xi1 forall a. Num a => a -> a -> a
* Double
xi1 forall a. Num a => a -> a -> a
+ Double
xi2 forall a. Num a => a -> a -> a
* Double
xi2
                        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef IORef Double
xi1Ref Double
xi1
                        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef IORef Double
xi2Ref Double
xi2
                        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef IORef Double
psiRef Double
psi
                        m ()
loop
                else forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef IORef Double
psiRef forall a b. (a -> b) -> a -> b
$ forall a. Floating a => a -> a
sqrt (- Double
2.0 forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
log Double
psi forall a. Fractional a => a -> a -> a
/ Double
psi)
     forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
       do Bool
flag <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO a
readIORef IORef Bool
flagRef
          if Bool
flag
            then do forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
flagRef Bool
False
                    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO a
readIORef IORef Double
nextRef
            else do forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef IORef Double
xi1Ref Double
0.0
                    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef IORef Double
xi2Ref Double
0.0
                    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef IORef Double
psiRef Double
0.0
                    m ()
loop
                    Double
xi1 <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO a
readIORef IORef Double
xi1Ref
                    Double
xi2 <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO a
readIORef IORef Double
xi2Ref
                    Double
psi <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO a
readIORef IORef Double
psiRef
                    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
flagRef Bool
True
                    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef IORef Double
nextRef forall a b. (a -> b) -> a -> b
$ Double
xi2 forall a. Num a => a -> a -> a
* Double
psi
                    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Double
xi1 forall a. Num a => a -> a -> a
* Double
psi