{-# LANGUAGE
        MultiParamTypeClasses,
        FlexibleInstances, FlexibleContexts,
        IncoherentInstances
  #-}

{-# OPTIONS_GHC -fno-warn-simplifiable-class-constraints #-}

module Data.Random.Sample where

import Control.Monad.State
import Control.Monad.Reader
import Data.Random.Distribution
import Data.Random.Lift
import Data.Random.RVar

import System.Random.Stateful

-- |A typeclass allowing 'Distribution's and 'RVar's to be sampled.  Both may
-- also be sampled via 'runRVar' or 'runRVarT', but I find it psychologically
-- pleasing to be able to sample both using this function, as they are two
-- separate abstractions for one base concept: a random variable.
class Sampleable d m t where
    -- |Directly sample from a distribution or random variable, using the given source of entropy.
    sampleFrom :: StatefulGen g m => g -> d t -> m t

instance Distribution d t => Sampleable d m t where
    sampleFrom :: forall g. StatefulGen g m => g -> d t -> m t
sampleFrom g
gen d t
d = forall (n :: * -> *) (m :: * -> *) g a.
(Lift n m, StatefulGen g m) =>
RVarT n a -> g -> m a
runRVarT (forall (d :: * -> *) t. Distribution d t => d t -> RVar t
rvar d t
d) g
gen

-- This instance overlaps with the other, but because RVarT is not a Distribution there is no conflict.
instance Lift m n => Sampleable (RVarT m) n t where
    sampleFrom :: forall g. StatefulGen g n => g -> RVarT m t -> n t
sampleFrom g
gen RVarT m t
x = forall (n :: * -> *) (m :: * -> *) g a.
(Lift n m, StatefulGen g m) =>
RVarT n a -> g -> m a
runRVarT RVarT m t
x g
gen

-- |Sample a random variable using the default source of entropy for the
-- monad in which the sampling occurs.
sample :: (Sampleable d m t, StatefulGen g m, MonadReader g m) => d t -> m t
sample :: forall (d :: * -> *) (m :: * -> *) t g.
(Sampleable d m t, StatefulGen g m, MonadReader g m) =>
d t -> m t
sample d t
thing = forall r (m :: * -> *). MonadReader r m => m r
ask forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \g
gen -> forall (d :: * -> *) (m :: * -> *) t g.
(Sampleable d m t, StatefulGen g m) =>
g -> d t -> m t
sampleFrom g
gen d t
thing

-- |Sample a random variable in a \"functional\" style.  Typical instantiations
-- of @s@ are @System.Random.StdGen@ or @System.Random.Mersenne.Pure64.PureMT@.
-- sample :: (Distribution d a, StatefulGen g m, MonadReader g m) => d t -> m t
-- sample thing gen = runStateGen gen (\stateGen -> sampleFrom stateGen thing)

sampleState :: (Distribution d t, RandomGen g, MonadState g m) => d t -> m t
sampleState :: forall (d :: * -> *) t g (m :: * -> *).
(Distribution d t, RandomGen g, MonadState g m) =>
d t -> m t
sampleState d t
thing = forall (d :: * -> *) (m :: * -> *) t g.
(Sampleable d m t, StatefulGen g m) =>
g -> d t -> m t
sampleFrom forall g. StateGenM g
StateGenM d t
thing

-- |Sample a random variable in a \"functional\" style.  Typical instantiations
-- of @g@ are @System.Random.StdGen@ or @System.Random.Mersenne.Pure64.PureMT@.
samplePure :: (Distribution d t, RandomGen g) => d t -> g -> (t, g)
samplePure :: forall (d :: * -> *) t g.
(Distribution d t, RandomGen g) =>
d t -> g -> (t, g)
samplePure d t
thing g
gen = forall g a.
RandomGen g =>
g -> (StateGenM g -> State g a) -> (a, g)
runStateGen g
gen (\StateGenM g
stateGen -> forall (d :: * -> *) (m :: * -> *) t g.
(Sampleable d m t, StatefulGen g m) =>
g -> d t -> m t
sampleFrom StateGenM g
stateGen d t
thing)