{-# LANGUAGE TypeFamilies #-}

-- |
-- Module     : Simulation.Aivika.RealTime.Generator
-- Copyright  : Copyright (c) 2016-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 'RT' monad can be an instance of 'MonadGenerator'.
--
module Simulation.Aivika.RealTime.Generator () where

import Control.Monad
import Control.Monad.Trans

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

import Data.IORef
import Data.Vector

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

import Simulation.Aivika.RealTime.Internal.RT

instance (Functor m, Monad m, MonadIO m) => MonadGenerator (RT m) where

  {-# SPECIALISE instance MonadGenerator (RT IO) #-}

  data Generator (RT m) =
    Generator { forall (m :: * -> *). Generator (RT m) -> RT m Double
generator01 :: RT m Double,
                -- ^ the generator of uniform numbers from 0 to 1
                forall (m :: * -> *). Generator (RT m) -> RT m Double
generatorNormal01 :: RT m Double,
                -- ^ the generator of normal numbers with mean 0 and variance 1
                forall (m :: * -> *). Generator (RT m) -> RT m Int
generatorSequenceNo :: RT m Int
                -- ^ the generator of sequence numbers
              }

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

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

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

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

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

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

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

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

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

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

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

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

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

  {-# INLINE generateSequenceNo #-}
  generateSequenceNo :: Generator (RT m) -> RT m Int
generateSequenceNo = forall (m :: * -> *). Generator (RT m) -> RT m Int
generatorSequenceNo

  {-# INLINABLE newGenerator #-}
  newGenerator :: GeneratorType (RT m) -> RT m (Generator (RT m))
newGenerator GeneratorType (RT m)
tp =
    case GeneratorType (RT m)
tp of
      GeneratorType (RT m)
SimpleGenerator ->
        do let g :: IO (IO Double)
g = 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
           IO Double
g' <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO (IO Double)
g
           forall (m :: * -> *).
MonadGenerator m =>
m Double -> m (Generator m)
newRandomGenerator01 (forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO Double
g')
      SimpleGeneratorWithSeed Word32
x ->
        do let g :: IO (IO Double)
g = 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)
           IO Double
g' <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO (IO Double)
g
           forall (m :: * -> *).
MonadGenerator m =>
m Double -> m (Generator m)
newRandomGenerator01 (forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO Double
g')
      CustomGenerator RT m (Generator (RT m))
g ->
        RT m (Generator (RT m))
g
      CustomGenerator01 RT m Double
g ->
        forall (m :: * -> *).
MonadGenerator m =>
m Double -> m (Generator m)
newRandomGenerator01 RT m Double
g

  {-# INLINABLE newRandomGenerator #-}
  newRandomGenerator :: forall g. RandomGen g => g -> RT m (Generator (RT m))
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 :: RT m 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 RT m Double
g01

  {-# INLINABLE newRandomGenerator01 #-}
  newRandomGenerator01 :: RT m Double -> RT m (Generator (RT m))
newRandomGenerator01 RT m Double
g01 =
    do RT m Double
gNormal01 <- forall (m :: * -> *). MonadIO m => m Double -> m (m Double)
newNormalGenerator01 RT m Double
g01
       IORef Int
gSeqNoRef <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. a -> IO (IORef a)
newIORef Int
0
       let gSeqNo :: RT m Int
gSeqNo =
             do Int
x <- 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 Int
gSeqNoRef
                forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ 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 :: RT m Double
generator01 = RT m Double
g01,
                          generatorNormal01 :: RT m Double
generatorNormal01 = RT m Double
gNormal01,
                          generatorSequenceNo :: RT m Int
generatorSequenceNo = RT m 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