module Foundation.Random.DRG
    ( RandomGen(..)
    , MonadRandomState(..)
    , withRandomGenerator
    ) where

import           Basement.Imports
import           Foundation.Random.Class

-- | A Deterministic Random Generator (DRG) class
class RandomGen gen where
    -- | Initialize a new random generator
    randomNew :: MonadRandom m => m gen

    -- | Initialize a new random generator from a binary seed.
    --
    -- If `Nothing` is returned, then the data is not acceptable
    -- for creating a new random generator.
    randomNewFrom :: UArray Word8 -> Maybe gen

    -- | Generate N bytes of randomness from a DRG
    randomGenerate :: CountOf Word8 -> gen -> (UArray Word8, gen)

    -- | Generate a Word64 from a DRG
    randomGenerateWord64 :: gen -> (Word64, gen)

    randomGenerateF32 :: gen -> (Float, gen)

    randomGenerateF64 :: gen -> (Double, gen)

-- | A simple Monad class very similar to a State Monad
-- with the state being a RandomGenerator.
newtype MonadRandomState gen a = MonadRandomState { MonadRandomState gen a -> gen -> (a, gen)
runRandomState :: gen -> (a, gen) }

instance Functor (MonadRandomState gen) where
    fmap :: (a -> b) -> MonadRandomState gen a -> MonadRandomState gen b
fmap a -> b
f MonadRandomState gen a
m = (gen -> (b, gen)) -> MonadRandomState gen b
forall gen a. (gen -> (a, gen)) -> MonadRandomState gen a
MonadRandomState ((gen -> (b, gen)) -> MonadRandomState gen b)
-> (gen -> (b, gen)) -> MonadRandomState gen b
forall a b. (a -> b) -> a -> b
$ \gen
g1 ->
        let (a
a, gen
g2) = MonadRandomState gen a -> gen -> (a, gen)
forall gen a. MonadRandomState gen a -> gen -> (a, gen)
runRandomState MonadRandomState gen a
m gen
g1 in (a -> b
f a
a, gen
g2)

instance Applicative (MonadRandomState gen) where
    pure :: a -> MonadRandomState gen a
pure a
a     = (gen -> (a, gen)) -> MonadRandomState gen a
forall gen a. (gen -> (a, gen)) -> MonadRandomState gen a
MonadRandomState ((gen -> (a, gen)) -> MonadRandomState gen a)
-> (gen -> (a, gen)) -> MonadRandomState gen a
forall a b. (a -> b) -> a -> b
$ \gen
g -> (a
a, gen
g)
    <*> :: MonadRandomState gen (a -> b)
-> MonadRandomState gen a -> MonadRandomState gen b
(<*>) MonadRandomState gen (a -> b)
fm MonadRandomState gen a
m = (gen -> (b, gen)) -> MonadRandomState gen b
forall gen a. (gen -> (a, gen)) -> MonadRandomState gen a
MonadRandomState ((gen -> (b, gen)) -> MonadRandomState gen b)
-> (gen -> (b, gen)) -> MonadRandomState gen b
forall a b. (a -> b) -> a -> b
$ \gen
g1 ->
        let (a -> b
f, gen
g2) = MonadRandomState gen (a -> b) -> gen -> (a -> b, gen)
forall gen a. MonadRandomState gen a -> gen -> (a, gen)
runRandomState MonadRandomState gen (a -> b)
fm gen
g1
            (a
a, gen
g3) = MonadRandomState gen a -> gen -> (a, gen)
forall gen a. MonadRandomState gen a -> gen -> (a, gen)
runRandomState MonadRandomState gen a
m gen
g2
         in (a -> b
f a
a, gen
g3)

instance Monad (MonadRandomState gen) where
    return :: a -> MonadRandomState gen a
return = a -> MonadRandomState gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    >>= :: MonadRandomState gen a
-> (a -> MonadRandomState gen b) -> MonadRandomState gen b
(>>=) MonadRandomState gen a
m1 a -> MonadRandomState gen b
m2 = (gen -> (b, gen)) -> MonadRandomState gen b
forall gen a. (gen -> (a, gen)) -> MonadRandomState gen a
MonadRandomState ((gen -> (b, gen)) -> MonadRandomState gen b)
-> (gen -> (b, gen)) -> MonadRandomState gen b
forall a b. (a -> b) -> a -> b
$ \gen
g1 ->
        let (a
a, gen
g2) = MonadRandomState gen a -> gen -> (a, gen)
forall gen a. MonadRandomState gen a -> gen -> (a, gen)
runRandomState MonadRandomState gen a
m1 gen
g1
         in MonadRandomState gen b -> gen -> (b, gen)
forall gen a. MonadRandomState gen a -> gen -> (a, gen)
runRandomState (a -> MonadRandomState gen b
m2 a
a) gen
g2

instance RandomGen gen => MonadRandom (MonadRandomState gen) where
    getRandomBytes :: CountOf Word8 -> MonadRandomState gen (UArray Word8)
getRandomBytes CountOf Word8
n = (gen -> (UArray Word8, gen)) -> MonadRandomState gen (UArray Word8)
forall gen a. (gen -> (a, gen)) -> MonadRandomState gen a
MonadRandomState (CountOf Word8 -> gen -> (UArray Word8, gen)
forall gen.
RandomGen gen =>
CountOf Word8 -> gen -> (UArray Word8, gen)
randomGenerate CountOf Word8
n)
    getRandomWord64 :: MonadRandomState gen Word64
getRandomWord64  = (gen -> (Word64, gen)) -> MonadRandomState gen Word64
forall gen a. (gen -> (a, gen)) -> MonadRandomState gen a
MonadRandomState gen -> (Word64, gen)
forall gen. RandomGen gen => gen -> (Word64, gen)
randomGenerateWord64
    getRandomF32 :: MonadRandomState gen Float
getRandomF32  = (gen -> (Float, gen)) -> MonadRandomState gen Float
forall gen a. (gen -> (a, gen)) -> MonadRandomState gen a
MonadRandomState gen -> (Float, gen)
forall gen. RandomGen gen => gen -> (Float, gen)
randomGenerateF32
    getRandomF64 :: MonadRandomState gen Double
getRandomF64  = (gen -> (Double, gen)) -> MonadRandomState gen Double
forall gen a. (gen -> (a, gen)) -> MonadRandomState gen a
MonadRandomState gen -> (Double, gen)
forall gen. RandomGen gen => gen -> (Double, gen)
randomGenerateF64


-- | Run a pure computation with a Random Generator in the 'MonadRandomState'
withRandomGenerator :: RandomGen gen
                    => gen
                    -> MonadRandomState gen a
                    -> (a, gen)
withRandomGenerator :: gen -> MonadRandomState gen a -> (a, gen)
withRandomGenerator gen
gen MonadRandomState gen a
m = MonadRandomState gen a -> gen -> (a, gen)
forall gen a. MonadRandomState gen a -> gen -> (a, gen)
runRandomState MonadRandomState gen a
m gen
gen