{-# LANGUAGE CPP                   #-}
{-# LANGUAGE TemplateHaskell       #-}
{-# LANGUAGE TypeApplications      #-}
{-|
Module      : Polysemy.RandomFu
Description : Polysemy random-fu effect

Polysemy "random-fu" effect.
This can be run in a few ways:
1. Directly in 'IO'
2. Using any 'Data.Random.RandomSource' from "random-fu"
3. In 'IO', using a given 'Data.Random.Source.PureMT' source.
('IO' is used to put the source in an 'IORef')

This module also contains the type-class instances to enable "absorbing"
MonadRandom, ala Polysemy.MTL.  See the tests for MTL or RandomFu for
examples of that in use.
-}

module Polysemy.RandomFu
  (
    -- * Effect
    RandomFu (..)

    -- * Actions
  , sampleRVar
#if MIN_VERSION_random_fu(0,3,0)
#else
  , getRandomPrim
#endif
  , sampleDist

    -- * Interpretations
  , runRandomSource
  , runRandomIO
  , runRandomIOPureMT
  )
where

import           Polysemy
import           Polysemy.State as PS
--import           Polysemy.MTL

import           Data.IORef                     ( newIORef )
import qualified Data.Random                   as R
import qualified Data.Random.RVar as R
import qualified Data.Random.Internal.Source   as R
import qualified Data.Random.Source.StdGen     as R
import qualified Data.Random.Source.PureMT     as R
import qualified Data.RVar                     as R (pureRVar)
import Data.Random.Source.IO ()
import           Control.Monad.IO.Class         ( MonadIO(..) )
import Control.Monad.Reader.Class (MonadReader)
import qualified System.Random.Stateful as SR
import GHC.IORef (IORef)

------------------------------------------------------------------------------
{- | An effect capable of sampling from a "random-fu" RVar or generating a
single random-variate of any type, @t@ with a
@Data.Random.Prim t@ constructor, currently one of @Word8@, @Word16@,
@Word32@, @Word64@, @Double@ or N-byte integer.
-}
data RandomFu m r where
  SampleRVar ::  R.RVar t -> RandomFu m t
#if MIN_VERSION_random_fu(0,3,0)
#else
  GetRandomPrim :: R.Prim t -> RandomFu m t
#endif
makeSem ''RandomFu

------------------------------------------------------------------------------
-- | use the 'RandomFu` effect to sample from a "random-fu" @Distribution@.
sampleDist
  :: (Member RandomFu r, R.Distribution d t) => d t -> Sem r t
sampleDist :: d t -> Sem r t
sampleDist = RVar t -> Sem r t
forall (r :: EffectRow) t. Member RandomFu r => RVar t -> Sem r t
sampleRVar (RVar t -> Sem r t) -> (d t -> RVar t) -> d t -> Sem r t
forall b c a. (b -> c) -> (a -> b) -> a -> c
. d t -> RVar t
forall (d :: * -> *) t. Distribution d t => d t -> RVar t
R.rvar
{-# INLINEABLE sampleDist #-}

------------------------------------------------------------------------------
-- | Run a 'Random' effect using a given 'R.RandomSource'
runRandomSource
  :: forall s m r a
   . (SR.StatefulGen s m
     , R.RandomSource m s
     , Member (Embed m) r
     )
  => s
  -> Sem (RandomFu ': r) a
  -> Sem r a
runRandomSource :: s -> Sem (RandomFu : r) a -> Sem r a
runRandomSource s
source = (forall (rInitial :: EffectRow) x.
 RandomFu (Sem rInitial) x -> Sem r x)
-> Sem (RandomFu : r) a -> Sem r a
forall (e :: (* -> *) -> * -> *) (r :: EffectRow) a.
FirstOrder e "interpret" =>
(forall (rInitial :: EffectRow) x. e (Sem rInitial) x -> Sem r x)
-> Sem (e : r) a -> Sem r a
interpret ((forall (rInitial :: EffectRow) x.
  RandomFu (Sem rInitial) x -> Sem r x)
 -> Sem (RandomFu : r) a -> Sem r a)
-> (forall (rInitial :: EffectRow) x.
    RandomFu (Sem rInitial) x -> Sem r x)
-> Sem (RandomFu : r) a
-> Sem r a
forall a b. (a -> b) -> a -> b
$ \case
    SampleRVar    rv -> m x -> Sem r x
forall (m :: * -> *) (r :: EffectRow) a.
Member (Embed m) r =>
m a -> Sem r a
embed (m x -> Sem r x) -> m x -> Sem r x
forall a b. (a -> b) -> a -> b
$ RVar x -> s -> m x
forall g (m :: * -> *) a. StatefulGen g m => RVar a -> g -> m a
R.runRVar RVar x
rv s
source
#if MIN_VERSION_random_fu(0,3,0)
#else
    GetRandomPrim pt -> embed $ R.getRandomPrimFrom source pt
#endif
{-# INLINEABLE runRandomSource #-}

------------------------------------------------------------------------------
-- | Run a 'Random` effect by using the default "random-fu" 'IO' source
runRandomIO
  :: forall r a
   . MonadIO (Sem r)
  => Sem (RandomFu ': r) a
  -> Sem r a
runRandomIO :: Sem (RandomFu : r) a -> Sem r a
runRandomIO Sem (RandomFu : r) a
x = do
  StdGen
g <- Sem r StdGen
forall (m :: * -> *). MonadIO m => m StdGen
SR.getStdGen
  StdGen -> Sem (State StdGen : r) a -> Sem r a
forall s (r :: EffectRow) a. s -> Sem (State s : r) a -> Sem r a
evalState StdGen
g (Sem (State StdGen : r) a -> Sem r a)
-> Sem (State StdGen : r) a -> Sem r a
forall a b. (a -> b) -> a -> b
$ (forall (rInitial :: EffectRow) x.
 RandomFu (Sem rInitial) x -> Sem (State StdGen : r) x)
-> Sem (RandomFu : r) a -> Sem (State StdGen : r) a
forall (e1 :: (* -> *) -> * -> *) (e2 :: (* -> *) -> * -> *)
       (r :: EffectRow) a.
FirstOrder e1 "reinterpret" =>
(forall (rInitial :: EffectRow) x.
 e1 (Sem rInitial) x -> Sem (e2 : r) x)
-> Sem (e1 : r) a -> Sem (e2 : r) a
reinterpret ( \case
    SampleRVar rv -> do
      StdGen
g <- Sem (State StdGen : r) StdGen
forall s (r :: EffectRow). Member (State s) r => Sem r s
get
      let (x
x, StdGen
g') = RVar x -> StdGen -> (x, StdGen)
forall g a. RandomGen g => RVar a -> g -> (a, g)
R.pureRVar RVar x
rv StdGen
g
      StdGen -> Sem (State StdGen : r) ()
forall s (r :: EffectRow). Member (State s) r => s -> Sem r ()
put StdGen
g'
      x -> Sem (State StdGen : r) x
forall (m :: * -> *) a. Monad m => a -> m a
return x
x
#if MIN_VERSION_random_fu(0,3,0)
#else
    GetRandomPrim pt -> raise $ R.getRandomPrim pt
#endif
    )
    Sem (RandomFu : r) a
x
{-# INLINEABLE runRandomIO #-}

------------------------------------------------------------------------------
-- | Run in 'IO', using the given 'R.PureMT' source, stored in an 'IORef'
runRandomIOPureMT
  :: (Member (Embed IO) r
      )
  => R.PureMT
  -> Sem (RandomFu ': r) a
  -> Sem r a
runRandomIOPureMT :: PureMT -> Sem (RandomFu : r) a -> Sem r a
runRandomIOPureMT PureMT
pMT = (forall (rInitial :: EffectRow) x.
 RandomFu (Sem rInitial) x -> Sem r x)
-> Sem (RandomFu : r) a -> Sem r a
forall (e :: (* -> *) -> * -> *) (r :: EffectRow) a.
FirstOrder e "interpret" =>
(forall (rInitial :: EffectRow) x. e (Sem rInitial) x -> Sem r x)
-> Sem (e : r) a -> Sem r a
interpret ((forall (rInitial :: EffectRow) x.
  RandomFu (Sem rInitial) x -> Sem r x)
 -> Sem (RandomFu : r) a -> Sem r a)
-> (forall (rInitial :: EffectRow) x.
    RandomFu (Sem rInitial) x -> Sem r x)
-> Sem (RandomFu : r) a
-> Sem r a
forall a b. (a -> b) -> a -> b
$ \case
   SampleRVar    rv -> IO x -> Sem r x
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO x -> Sem r x) -> IO x -> Sem r x
forall a b. (a -> b) -> a -> b
$ do
     IOGenM PureMT
g <- PureMT -> IO (IOGenM PureMT)
forall (m :: * -> *) g. MonadIO m => g -> m (IOGenM g)
SR.newIOGenM PureMT
pMT
     RVar x -> IOGenM PureMT -> IO x
forall g (m :: * -> *) a. StatefulGen g m => RVar a -> g -> m a
R.runRVar RVar x
rv IOGenM PureMT
g
#if MIN_VERSION_random_fu(0,3,0)
#else
   GetRandomPrim pt -> liftIO $ do
     g <- newIORef pMT
     R.getRandomPrimFromMTRef g pt
#endif
--getRandomPrimFromMTRef   embed (SR.newIOGenM source) >>= flip runRandomSource re
{-# INLINEABLE runRandomIOPureMT #-}