{-# LANGUAGE TypeFamilies, FlexibleInstances #-}

-- |
-- Module     : Simulation.Aivika.Branch.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 7.10.3
--
-- Here is defined a random number generator,
-- where 'BR' can be an instance of 'MonadGenerator'.
--
module Simulation.Aivika.Branch.Generator () where

import Control.Monad
import Control.Monad.Trans

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

import Data.IORef

import Simulation.Aivika.Trans
import Simulation.Aivika.Trans.Generator.Primitive
import Simulation.Aivika.Branch.Internal.BR

instance MonadGenerator (BR IO) where

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

  generateUniform :: Generator (BR IO) -> Double -> Double -> BR 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 (BR IO) -> BR IO Double
generator01

  generateUniformInt :: Generator (BR IO) -> Int -> Int -> BR 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 (BR IO) -> BR IO Double
generator01

  generateTriangular :: Generator (BR IO) -> Double -> Double -> Double -> BR 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 (BR IO) -> BR IO Double
generator01

  generateNormal :: Generator (BR IO) -> Double -> Double -> BR 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 (BR IO) -> BR IO Double
generatorNormal01

  generateLogNormal :: Generator (BR IO) -> Double -> Double -> BR 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 (BR IO) -> BR IO Double
generatorNormal01

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

  generateErlang :: Generator (BR IO) -> Double -> Int -> BR 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 (BR IO) -> BR IO Double
generator01

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

  generateBinomial :: Generator (BR IO) -> Double -> Int -> BR 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 (BR IO) -> BR IO Double
generator01

  generateGamma :: Generator (BR IO) -> Double -> Double -> BR IO Double
generateGamma Generator (BR IO)
g = forall (m :: * -> *).
Monad m =>
m Double -> m Double -> Double -> Double -> m Double
generateGamma01 (Generator (BR IO) -> BR IO Double
generatorNormal01 Generator (BR IO)
g) (Generator (BR IO) -> BR IO Double
generator01 Generator (BR IO)
g)

  generateBeta :: Generator (BR IO) -> Double -> Double -> BR IO Double
generateBeta Generator (BR IO)
g = forall (m :: * -> *).
Monad m =>
m Double -> m Double -> Double -> Double -> m Double
generateBeta01 (Generator (BR IO) -> BR IO Double
generatorNormal01 Generator (BR IO)
g) (Generator (BR IO) -> BR IO Double
generator01 Generator (BR IO)
g)

  generateWeibull :: Generator (BR IO) -> Double -> Double -> BR 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 (BR IO) -> BR IO Double
generator01

  generateDiscrete :: forall a. Generator (BR IO) -> DiscretePDF a -> BR 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 (BR IO) -> BR IO Double
generator01

  generateSequenceNo :: Generator (BR IO) -> BR IO Int
generateSequenceNo = Generator (BR IO) -> BR IO Int
generatorSequenceNo

  newGenerator :: GeneratorType (BR IO) -> BR IO (Generator (BR IO))
newGenerator GeneratorType (BR IO)
tp =
    case GeneratorType (BR IO)
tp of
      GeneratorType (BR IO)
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 ->
        forall a. HasCallStack => [Char] -> a
error [Char]
"Unsupported generator type SimpleGeneratorWithSeed: newGenerator"
      CustomGenerator BR IO (Generator (BR IO))
g ->
        BR IO (Generator (BR IO))
g
      CustomGenerator01 BR IO Double
g ->
        forall (m :: * -> *).
MonadGenerator m =>
m Double -> m (Generator m)
newRandomGenerator01 BR IO Double
g

  newRandomGenerator :: forall g. RandomGen g => g -> BR IO (Generator (BR 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 :: BR 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 BR IO Double
g01

  newRandomGenerator01 :: BR IO Double -> BR IO (Generator (BR IO))
newRandomGenerator01 BR IO Double
g01 =
    do BR IO Double
gNormal01 <- BR IO Double -> BR IO (BR IO Double)
newNormalGenerator01 BR IO 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 :: BR IO 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 :: BR IO Double
generator01 = BR IO Double
g01,
                          generatorNormal01 :: BR IO Double
generatorNormal01 = BR IO Double
gNormal01,
                          generatorSequenceNo :: BR IO Int
generatorSequenceNo = BR 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 :: BR IO Double
                        -- ^ the generator
                        -> BR IO (BR IO Double)
newNormalGenerator01 :: BR IO Double -> BR IO (BR IO Double)
newNormalGenerator01 BR IO 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 :: BR IO ()
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 <- BR IO Double
g
                        Double
g2 <- BR IO 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
                        BR IO ()
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
                    BR IO ()
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