-- | A monad that abstract randomized operations.
module MCSP.System.Random.Monad (
    Random,
    evalRandom,
    liftRandom,
    lazyRandom,
) where

import Control.Applicative (Applicative (..))
import Control.Monad (Monad (..))
import Control.Monad.ST (ST)
import Control.Monad.ST.Unsafe (unsafeInterleaveST)
import Control.Monad.Trans.Reader (ReaderT (..), mapReaderT, runReaderT)
import Data.Function (($), (.))
import Data.Functor (Functor (..), (<$>))
import Data.Monoid (Monoid (..))
import Data.Semigroup (Semigroup (..))
import Data.Traversable (sequence)
import System.Random.PCG.Class (Generator)

-- ------------ --
-- Random Monad --
-- ------------ --

-- | A monad capable of producing random values of @a@.
newtype Random a = Random (forall g s. Generator g (ST s) => ReaderT g (ST s) a)

instance Functor Random where
    fmap :: forall a b. (a -> b) -> Random a -> Random b
fmap a -> b
f (Random forall g s. Generator g (ST s) => ReaderT g (ST s) a
gen) = (forall g s. Generator g (ST s) => ReaderT g (ST s) b) -> Random b
forall a.
(forall g s. Generator g (ST s) => ReaderT g (ST s) a) -> Random a
Random ((a -> b) -> ReaderT g (ST s) a -> ReaderT g (ST s) b
forall a b. (a -> b) -> ReaderT g (ST s) a -> ReaderT g (ST s) b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f ReaderT g (ST s) a
forall g s. Generator g (ST s) => ReaderT g (ST s) a
gen)
    {-# INLINE fmap #-}
    a
x <$ :: forall a b. a -> Random b -> Random a
<$ Random b
_ = a -> Random a
forall a. a -> Random a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x
    {-# INLINE (<$) #-}

instance Applicative Random where
    pure :: forall a. a -> Random a
pure a
x = (forall g s. Generator g (ST s) => ReaderT g (ST s) a) -> Random a
forall a.
(forall g s. Generator g (ST s) => ReaderT g (ST s) a) -> Random a
Random (a -> ReaderT g (ST s) a
forall a. a -> ReaderT g (ST s) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x)
    {-# INLINE pure #-}
    liftA2 :: forall a b c. (a -> b -> c) -> Random a -> Random b -> Random c
liftA2 a -> b -> c
f (Random forall g s. Generator g (ST s) => ReaderT g (ST s) a
genA) (Random forall g s. Generator g (ST s) => ReaderT g (ST s) b
genB) = (forall g s. Generator g (ST s) => ReaderT g (ST s) c) -> Random c
forall a.
(forall g s. Generator g (ST s) => ReaderT g (ST s) a) -> Random a
Random ((a -> b -> c)
-> ReaderT g (ST s) a -> ReaderT g (ST s) b -> ReaderT g (ST s) c
forall a b c.
(a -> b -> c)
-> ReaderT g (ST s) a -> ReaderT g (ST s) b -> ReaderT g (ST s) c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> b -> c
f ReaderT g (ST s) a
forall g s. Generator g (ST s) => ReaderT g (ST s) a
genA ReaderT g (ST s) b
forall g s. Generator g (ST s) => ReaderT g (ST s) b
genB)
    {-# INLINE liftA2 #-}
    Random forall g s. Generator g (ST s) => ReaderT g (ST s) (a -> b)
genF <*> :: forall a b. Random (a -> b) -> Random a -> Random b
<*> Random forall g s. Generator g (ST s) => ReaderT g (ST s) a
genA = (forall g s. Generator g (ST s) => ReaderT g (ST s) b) -> Random b
forall a.
(forall g s. Generator g (ST s) => ReaderT g (ST s) a) -> Random a
Random (ReaderT g (ST s) (a -> b)
forall g s. Generator g (ST s) => ReaderT g (ST s) (a -> b)
genF ReaderT g (ST s) (a -> b)
-> ReaderT g (ST s) a -> ReaderT g (ST s) b
forall a b.
ReaderT g (ST s) (a -> b)
-> ReaderT g (ST s) a -> ReaderT g (ST s) b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ReaderT g (ST s) a
forall g s. Generator g (ST s) => ReaderT g (ST s) a
genA)
    {-# INLINE (<*>) #-}
    Random forall g s. Generator g (ST s) => ReaderT g (ST s) a
genA *> :: forall a b. Random a -> Random b -> Random b
*> Random forall g s. Generator g (ST s) => ReaderT g (ST s) b
genB = (forall g s. Generator g (ST s) => ReaderT g (ST s) b) -> Random b
forall a.
(forall g s. Generator g (ST s) => ReaderT g (ST s) a) -> Random a
Random (ReaderT g (ST s) a
forall g s. Generator g (ST s) => ReaderT g (ST s) a
genA ReaderT g (ST s) a -> ReaderT g (ST s) b -> ReaderT g (ST s) b
forall a b.
ReaderT g (ST s) a -> ReaderT g (ST s) b -> ReaderT g (ST s) b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> ReaderT g (ST s) b
forall g s. Generator g (ST s) => ReaderT g (ST s) b
genB)
    {-# INLINE (*>) #-}
    Random forall g s. Generator g (ST s) => ReaderT g (ST s) a
genA <* :: forall a b. Random a -> Random b -> Random a
<* Random forall g s. Generator g (ST s) => ReaderT g (ST s) b
genB = (forall g s. Generator g (ST s) => ReaderT g (ST s) a) -> Random a
forall a.
(forall g s. Generator g (ST s) => ReaderT g (ST s) a) -> Random a
Random (ReaderT g (ST s) a
forall g s. Generator g (ST s) => ReaderT g (ST s) a
genA ReaderT g (ST s) a -> ReaderT g (ST s) b -> ReaderT g (ST s) a
forall a b.
ReaderT g (ST s) a -> ReaderT g (ST s) b -> ReaderT g (ST s) a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* ReaderT g (ST s) b
forall g s. Generator g (ST s) => ReaderT g (ST s) b
genB)
    {-# INLINE (<*) #-}

instance Monad Random where
    Random forall g s. Generator g (ST s) => ReaderT g (ST s) a
genA >>= :: forall a b. Random a -> (a -> Random b) -> Random b
>>= a -> Random b
f = (forall g s. Generator g (ST s) => ReaderT g (ST s) b) -> Random b
forall a.
(forall g s. Generator g (ST s) => ReaderT g (ST s) a) -> Random a
Random (ReaderT g (ST s) a
forall g s. Generator g (ST s) => ReaderT g (ST s) a
genA ReaderT g (ST s) a
-> (a -> ReaderT g (ST s) b) -> ReaderT g (ST s) b
forall a b.
ReaderT g (ST s) a
-> (a -> ReaderT g (ST s) b) -> ReaderT g (ST s) b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Random b -> ReaderT g (ST s) b
forall {g} {s} {a}.
Generator g (ST s) =>
Random a -> ReaderT g (ST s) a
getRandom (Random b -> ReaderT g (ST s) b)
-> (a -> Random b) -> a -> ReaderT g (ST s) b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Random b
f)
      where
        getRandom :: Random a -> ReaderT g (ST s) a
getRandom (Random forall g s. Generator g (ST s) => ReaderT g (ST s) a
gen) = ReaderT g (ST s) a
forall g s. Generator g (ST s) => ReaderT g (ST s) a
gen
    {-# INLINE (>>=) #-}

instance Semigroup a => Semigroup (Random a) where
    <> :: Random a -> Random a -> Random a
(<>) = (a -> a -> a) -> Random a -> Random a -> Random a
forall a b c. (a -> b -> c) -> Random a -> Random b -> Random c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> a -> a
forall a. Semigroup a => a -> a -> a
(<>)
    {-# INLINE (<>) #-}
    sconcat :: NonEmpty (Random a) -> Random a
sconcat NonEmpty (Random a)
xs = NonEmpty a -> a
forall a. Semigroup a => NonEmpty a -> a
sconcat (NonEmpty a -> a) -> Random (NonEmpty a) -> Random a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NonEmpty (Random a) -> Random (NonEmpty a)
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => NonEmpty (m a) -> m (NonEmpty a)
sequence NonEmpty (Random a)
xs
    {-# INLINE sconcat #-}
    stimes :: forall b. Integral b => b -> Random a -> Random a
stimes b
n Random a
x = b -> a -> a
forall b. Integral b => b -> a -> a
forall a b. (Semigroup a, Integral b) => b -> a -> a
stimes b
n (a -> a) -> Random a -> Random a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Random a
x
    {-# INLINE stimes #-}

instance Monoid a => Monoid (Random a) where
    mempty :: Random a
mempty = a -> Random a
forall a. a -> Random a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
forall a. Monoid a => a
mempty
    {-# INLINE mempty #-}
    mconcat :: [Random a] -> Random a
mconcat [Random a]
xs = [a] -> a
forall a. Monoid a => [a] -> a
mconcat ([a] -> a) -> Random [a] -> Random a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Random a] -> Random [a]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence [Random a]
xs
    {-# INLINE mconcat #-}

-- ---------- --
-- Evaluation --
-- ---------- --

-- | Evaluate a random computation with the given initial generator and return the final value.
evalRandom :: Generator g (ST s) => Random a -> g -> ST s a
evalRandom :: forall g s a. Generator g (ST s) => Random a -> g -> ST s a
evalRandom (Random forall g s. Generator g (ST s) => ReaderT g (ST s) a
gen) = ReaderT g (ST s) a -> g -> ST s a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT g (ST s) a
forall g s. Generator g (ST s) => ReaderT g (ST s) a
gen
{-# INLINE evalRandom #-}

-- | Turn a standard RNG function into a `Random` monad.
liftRandom :: (forall g m. Generator g m => g -> m a) -> Random a
liftRandom :: forall a.
(forall g (m :: * -> *). Generator g m => g -> m a) -> Random a
liftRandom forall g (m :: * -> *). Generator g m => g -> m a
gen = (forall g s. Generator g (ST s) => ReaderT g (ST s) a) -> Random a
forall a.
(forall g s. Generator g (ST s) => ReaderT g (ST s) a) -> Random a
Random ((forall g s. Generator g (ST s) => ReaderT g (ST s) a)
 -> Random a)
-> (forall g s. Generator g (ST s) => ReaderT g (ST s) a)
-> Random a
forall a b. (a -> b) -> a -> b
$ (g -> ST s a) -> ReaderT g (ST s) a
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT g -> ST s a
forall g (m :: * -> *). Generator g m => g -> m a
gen
{-# INLINE liftRandom #-}

-- | Allows a `Random` monad to be evaluated lazily.
--
-- This function should be applied with care, otherwise a strict `Random` monad could still force
-- early evaluation of this lazy version.
lazyRandom :: Random a -> Random a
lazyRandom :: forall a. Random a -> Random a
lazyRandom (Random forall g s. Generator g (ST s) => ReaderT g (ST s) a
gen) =
    -- SAFETY: unsafeInterleaveST break the order of ST operations, which shouldn't matter for
    -- random operations (except of changing the output value)
    (forall g s. Generator g (ST s) => ReaderT g (ST s) a) -> Random a
forall a.
(forall g s. Generator g (ST s) => ReaderT g (ST s) a) -> Random a
Random ((forall g s. Generator g (ST s) => ReaderT g (ST s) a)
 -> Random a)
-> (forall g s. Generator g (ST s) => ReaderT g (ST s) a)
-> Random a
forall a b. (a -> b) -> a -> b
$ (ST s a -> ST s a) -> ReaderT g (ST s) a -> ReaderT g (ST s) a
forall (m :: * -> *) a (n :: * -> *) b r.
(m a -> n b) -> ReaderT r m a -> ReaderT r n b
mapReaderT ST s a -> ST s a
forall s a. ST s a -> ST s a
unsafeInterleaveST ReaderT g (ST s) a
forall g s. Generator g (ST s) => ReaderT g (ST s) a
gen
{-# INLINE lazyRandom #-}