{-# LANGUAGE ScopedTypeVariables #-}

-- |
-- Module      :  Mcmc.Internal.Random
-- Description :  Tools for random calculations
-- Copyright   :  (c) Dominik Schrempf, 2021
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Wed Nov 25 07:14:52 2020.
module Mcmc.Internal.Random
  ( splitGen,
    saveGen,
    loadGen,
  )
where

import Control.Monad
import Control.Monad.Primitive
import qualified Data.Vector.Unboxed as V
import Data.Word
import System.Random.MWC

-- | Split a generator.
--
-- Splitting an MWC pseudo number generator is not good practice. However, I
-- have to go with this solution for now, and wait for proper support of
-- spittable pseudo random number generators such as @splitmix@.
splitGen :: PrimMonad m => Int -> Gen (PrimState m) -> m [Gen (PrimState m)]
splitGen :: Int -> Gen (PrimState m) -> m [Gen (PrimState m)]
splitGen Int
n Gen (PrimState m)
gen
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = [Gen (PrimState m)] -> m [Gen (PrimState m)]
forall (m :: * -> *) a. Monad m => a -> m a
return []
  | Bool
otherwise = do
      [Vector Word32]
seeds :: [V.Vector Word32] <- Int -> m (Vector Word32) -> m [Vector Word32]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n (m (Vector Word32) -> m [Vector Word32])
-> m (Vector Word32) -> m [Vector Word32]
forall a b. (a -> b) -> a -> b
$ Gen (PrimState m) -> Int -> m (Vector Word32)
forall (m :: * -> *) g a (v :: * -> *).
(PrimMonad m, StatefulGen g m, Uniform a, Vector v a) =>
g -> Int -> m (v a)
uniformVector Gen (PrimState m)
gen Int
256
      (Vector Word32 -> m (Gen (PrimState m)))
-> [Vector Word32] -> m [Gen (PrimState m)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Vector Word32 -> m (Gen (PrimState m))
forall (m :: * -> *) (v :: * -> *).
(PrimMonad m, Vector v Word32) =>
v Word32 -> m (Gen (PrimState m))
initialize [Vector Word32]
seeds

-- TODO: Splitmix. Remove or amend these functions as soon as split mix is used
-- and is available with the statistics package.

-- | Save a generator to a seed.
saveGen :: GenIO -> IO (V.Vector Word32)
saveGen :: GenIO -> IO (Vector Word32)
saveGen = (Seed -> Vector Word32) -> IO Seed -> IO (Vector Word32)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Seed -> Vector Word32
fromSeed (IO Seed -> IO (Vector Word32))
-> (Gen RealWorld -> IO Seed)
-> Gen RealWorld
-> IO (Vector Word32)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Gen RealWorld -> IO Seed
forall (m :: * -> *). PrimMonad m => Gen (PrimState m) -> m Seed
save

-- | Load a generator from a seed.
loadGen :: V.Vector Word32 -> IO GenIO
loadGen :: Vector Word32 -> IO GenIO
loadGen = Seed -> IO (Gen RealWorld)
forall (m :: * -> *). PrimMonad m => Seed -> m (Gen (PrimState m))
restore (Seed -> IO (Gen RealWorld))
-> (Vector Word32 -> Seed) -> Vector Word32 -> IO (Gen RealWorld)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Word32 -> Seed
forall (v :: * -> *). Vector v Word32 => v Word32 -> Seed
toSeed