content: 
  {-# LANGUAGE RankNTypes #-}
  
  -- |
  -- Module     : Simulation.Aivika.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
  --
  -- Below is defined a type class of the random number generator.
  --
  module Simulation.Aivika.Generator 
         (Generator(..),
          GeneratorType(..),
          DiscretePDF(..),
          newGenerator,
          newRandomGenerator,
          newRandomGenerator01) where
  
  import System.Random
  import qualified System.Random.MWC as MWC
  import Data.IORef
  import Data.Word
  import Data.Vector
  import Data.Functor
  
  -- | A discrete probability density function.
  type DiscretePDF a = [(a, Double)]
  
  -- | Defines a random number generator.
  data Generator =
    Generator { generateUniform :: Double -> Double -> IO Double,
                -- ^ Generate an uniform random number
                -- with the specified minimum and maximum.
                generateUniformInt :: Int -> Int -> IO Int,
                -- ^ Generate an uniform integer random number
                -- with the specified minimum and maximum.
                generateTriangular :: Double -> Double -> Double -> IO Double,
                -- ^ Generate a triangular random number
                -- by the specified minimum, median and maximum.
                generateNormal :: Double -> Double -> IO Double,
                -- ^ Generate the normal random number
                -- with the specified mean and deviation.
                generateLogNormal :: Double -> Double -> IO Double,
                -- ^ Generate a random number from the lognormal distribution derived
                -- from a normal distribution with the specified mean and deviation.
                generateExponential :: Double -> IO Double,
                -- ^ Generate the random number distributed exponentially
                -- with the specified mean (the reciprocal of the rate).
                generateErlang :: Double -> Int -> IO Double,
                -- ^ Generate the Erlang random number
                -- with the specified scale (the reciprocal of the rate) and integer shape.
                generatePoisson :: Double -> IO Int,
                -- ^ Generate the Poisson random number
                -- with the specified mean.
                generateBinomial :: Double -> Int -> IO Int,
                -- ^ Generate the binomial random number
                -- with the specified probability and number of trials.
                generateGamma :: Double -> Double -> IO Double,
                -- ^ Generate a random number from the Gamma distribution with
                -- the specified shape (kappa) and scale (theta, a reciprocal of the rate).
                --
                -- The probability density for the Gamma distribution is
                --
                -- @f x = x ** (kappa - 1) * exp (- x \/ theta) \/ theta ** kappa * Gamma kappa@
                generateBeta :: Double -> Double -> IO Double,
                -- ^ Generate a random number from the Beta distribution by
                -- the specified shape parameters (alpha and beta).
                --
                -- The probability density for the Beta distribution is
                --
                -- @f x = x ** (alpha - 1) * (1 - x) ** (beta - 1) \/ B alpha beta@
                generateWeibull :: Double -> Double -> IO Double,
                -- ^ Generate a random number from the Weibull distribution by
                -- the specified shape and scale.
                generateDiscrete :: forall a. DiscretePDF a -> IO a,
                -- ^ Generate a random value from the specified discrete distribution.
                generateSequenceNo :: IO Int
                -- ^ Generate a sequence number which can be considered quite unique.
              }
  
  -- | Generate the uniform random number with the specified minimum and maximum.
  generateUniform01 :: IO Double
                       -- ^ the generator
                       -> Double
                       -- ^ minimum
                       -> Double
                       -- ^ maximum
                       -> IO Double
  generateUniform01 g min max =
    do x <- g
       return $ min + x * (max - min)
  
  -- | Generate the uniform random number with the specified minimum and maximum.
  generateUniformInt01 :: IO Double
                          -- ^ the generator
                          -> Int
                          -- ^ minimum
                          -> Int
                          -- ^ maximum
                          -> IO Int
  generateUniformInt01 g min max =
    do x <- g
       let min' = fromIntegral min - 0.5
           max' = fromIntegral max + 0.5
           z    = round (min' + x * (max' - min'))
           z'   = if z < min
                  then min
                  else if z > max
                       then max
                       else z
       return z'
  
  -- | Generate the triangular random number by the specified minimum, median and maximum.
  generateTriangular01 :: IO Double
                          -- ^ the generator
                          -> Double
                          -- ^ minimum
                          -> Double
                          -- ^ median
                          -> Double
                          -- ^ maximum
                          -> IO Double
  generateTriangular01 g min median max =
    do x <- g
       if x <= (median - min) / (max - min)
         then return $ min + sqrt ((median - min) * (max - min) * x)
         else return $ max - sqrt ((max - median) * (max - min) * (1 - x))
  
  -- | 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 :: IO Double
                          -- ^ the generator
                          -> IO (IO Double)
  newNormalGenerator01 g =
    do nextRef <- newIORef 0.0
       flagRef <- newIORef False
       xi1Ref  <- newIORef 0.0
       xi2Ref  <- newIORef 0.0
       psiRef  <- newIORef 0.0
       let loop =
             do psi <- readIORef psiRef
                if (psi >= 1.0) || (psi == 0.0)
                  then do g1 <- g
                          g2 <- g
                          let xi1 = 2.0 * g1 - 1.0
                              xi2 = 2.0 * g2 - 1.0
                              psi = xi1 * xi1 + xi2 * xi2
                          writeIORef xi1Ref xi1
                          writeIORef xi2Ref xi2
                          writeIORef psiRef psi
                          loop
                  else writeIORef psiRef $ sqrt (- 2.0 * log psi / psi)
       return $
         do flag <- readIORef flagRef
            if flag
              then do writeIORef flagRef False
                      readIORef nextRef
              else do writeIORef xi1Ref 0.0
                      writeIORef xi2Ref 0.0
                      writeIORef psiRef 0.0
                      loop
                      xi1 <- readIORef xi1Ref
                      xi2 <- readIORef xi2Ref
                      psi <- readIORef psiRef
                      writeIORef flagRef True
                      writeIORef nextRef $ xi2 * psi
                      return $ xi1 * psi
  
  -- | Return the exponential random number with the specified mean.
  generateExponential01 :: IO Double
                           -- ^ the generator
                           -> Double
                           -- ^ the mean
                           -> IO Double
  generateExponential01 g mu =
    do x <- g
       return (- log x * mu)
  
  -- | Return the Erlang random number.
  generateErlang01 :: IO Double
                      -- ^ the generator
                      -> Double
                      -- ^ the scale
                      -> Int
                      -- ^ the shape
                      -> IO Double
  generateErlang01 g beta m =
    do x <- loop m 1
       return (- log x * beta)
         where loop m acc
                 | m < 0     = error "Negative shape: generateErlang."
                 | m == 0    = return acc
                 | otherwise = do x <- g
                                  loop (m - 1) (x * acc)
  
  -- | Generate the Poisson random number with the specified mean.
  generatePoisson01 :: IO Double
                       -- ^ the generator
                       -> Double
                       -- ^ the mean
                       -> IO Int
  generatePoisson01 g mu =
    do prob0 <- g
       let loop prob prod acc
             | prob <= prod = return acc
             | otherwise    = loop
                              (prob - prod)
                              (prod * mu / fromIntegral (acc + 1))
                              (acc + 1)
       loop prob0 (exp (- mu)) 0
  
  -- | Generate a binomial random number with the specified probability and number of trials. 
  generateBinomial01 :: IO Double
                        -- ^ the generator
                        -> Double 
                        -- ^ the probability
                        -> Int
                        -- ^ the number of trials
                        -> IO Int
  generateBinomial01 g prob trials = loop trials 0 where
    loop n acc
      | n < 0     = error "Negative number of trials: generateBinomial."
      | n == 0    = return acc
      | otherwise = do x <- g
                       if x <= prob
                         then loop (n - 1) (acc + 1)
                         else loop (n - 1) acc
  
  -- | Generate a random number from the Gamma distribution using Marsaglia and Tsang method.
  generateGamma01 :: IO Double
                     -- ^ a normal random number ~ N (0,1)
                     -> IO Double
                     -- ^ an uniform random number ~ U (0, 1)
                     -> Double
                     -- ^ the shape parameter (kappa) 
                     -> Double
                     -- ^ the scale parameter (theta)
                     -> IO Double
  generateGamma01 gn gu kappa theta
    | kappa <= 0 = error "The shape parameter (kappa) must be positive: generateGamma01"
    | kappa > 1  =
      let d = kappa - 1 / 3
          c = 1 / sqrt (9 * d)
          loop =
            do z <- gn
               if z <= - (1 / c)
                 then loop
                 else do let v = (1 + c * z) ** 3
                         u <- gu
                         if log u > 0.5 * z * z + d - d * v + d * log v
                           then loop
                           else return $ d * v * theta
      in loop
    | otherwise  =
      do x <- generateGamma01 gn gu (1 + kappa) theta
         u <- gu
         return $ x * u ** (1 / kappa)
  
  -- | Generate a random number from the Beta distribution.
  generateBeta01 :: IO Double
                    -- ^ a normal random number ~ N (0, 1)
                    -> IO Double
                    -- ^ an uniform random number ~ U (0, 1)
                    -> Double
                    -- ^ the shape parameter alpha
                    -> Double
                    -- ^ the shape parameter beta
                    -> IO Double
  generateBeta01 gn gu alpha beta =
    do g1 <- generateGamma01 gn gu alpha 1
       g2 <- generateGamma01 gn gu beta 1
       return $ g1 / (g1 + g2)
  
  -- | Generate a random number from the Weibull distribution.
  generateWeibull01 :: IO Double
                       -- ^ an uniform random number ~ U (0, 1)
                       -> Double
                       -- ^ shape
                       -> Double
                       -- ^ scale
                       -> IO Double
  generateWeibull01 g alpha beta =
    do x <- g
       return $ beta * (- log x) ** (1 / alpha)
  
  -- | Generate a random value from the specified discrete distribution.
  generateDiscrete01 :: IO Double
                        -- ^ an uniform random number ~ U (0, 1)
                        -> DiscretePDF a
                        -- ^ a discrete probability density function
                        -> IO a
  generateDiscrete01 g []   = error "Empty PDF: generateDiscrete01"
  generateDiscrete01 g dpdf =
    do x <- g
       let loop acc [(a, p)] = a
           loop acc ((a, p) : dpdf) =
             if x <= acc + p
             then a
             else loop (acc + p) dpdf
       return $ loop 0 dpdf
  
  -- | Defines a type of the random number generator.
  data GeneratorType = SimpleGenerator
                       -- ^ The simple random number generator.
                     | SimpleGeneratorWithSeed Word32
                       -- ^ The simple random number generator with the specified seed.
                     | CustomGenerator (IO Generator)
                       -- ^ The custom random number generator.
                     | CustomGenerator01 (IO Double)
                       -- ^ The custom random number generator by the specified uniform
                       -- generator of numbers from 0 to 1.
  
  -- | Create a new random number generator by the specified type.
  newGenerator :: GeneratorType -> IO Generator
  newGenerator tp =
    case tp of
      SimpleGenerator ->
        MWC.uniform <$> MWC.createSystemRandom >>= newRandomGenerator01
      SimpleGeneratorWithSeed x ->
        MWC.uniform <$> MWC.initialize (singleton x) >>= newRandomGenerator01
      CustomGenerator g ->
        g
      CustomGenerator01 g ->
        newRandomGenerator01 g
  
  -- | Create a new random generator by the specified standard generator.
  newRandomGenerator :: RandomGen g => g -> IO Generator
  newRandomGenerator g =
    do r <- newIORef g
       let g1 = do g <- readIORef r
                   let (x, g') = random g
                   writeIORef r g'
                   return x
       newRandomGenerator01 g1
  
  -- | Create a new random generator by the specified uniform generator of numbers from 0 to 1.
  newRandomGenerator01 :: IO Double -> IO Generator
  newRandomGenerator01 g =
    do let g1 = g
       g2 <- newNormalGenerator01 g1
       let g3 mu nu = do { x <- g2; return $ mu + nu * x }
           g4 mu nu = do { x <- g2; return $ exp (mu + nu * x) }
       seqNoRef <- newIORef 0
       let seqNo = do { x <- readIORef seqNoRef; modifyIORef' seqNoRef (+1); return x }
       return Generator { generateUniform = generateUniform01 g1,
                          generateUniformInt = generateUniformInt01 g1,
                          generateTriangular = generateTriangular01 g1,
                          generateNormal = g3,
                          generateLogNormal = g4,
                          generateExponential = generateExponential01 g1,
                          generateErlang = generateErlang01 g1,
                          generatePoisson = generatePoisson01 g1,
                          generateBinomial = generateBinomial01 g1,
                          generateGamma = generateGamma01 g2 g1,
                          generateBeta = generateBeta01 g2 g1,
                          generateWeibull = generateWeibull01 g1,
                          generateDiscrete = generateDiscrete01 g1,
                          generateSequenceNo = seqNo }
  