{-# LANGUAGE TupleSections #-}
module Lighthouse.Utils.Random
    ( RandomM (..)
    , runRandomMIO, getRandomGen, randomM, randomRM, nRandomsR, nRandoms
    ) where

import System.Random

newtype RandomM g a = RandomM { forall g a. RandomM g a -> g -> (a, g)
runRandomM :: g -> (a, g) }

instance Functor (RandomM g) where
    fmap :: forall a b. (a -> b) -> RandomM g a -> RandomM g b
fmap a -> b
f RandomM g a
r = (g -> (b, g)) -> RandomM g b
forall g a. (g -> (a, g)) -> RandomM g a
RandomM ((g -> (b, g)) -> RandomM g b) -> (g -> (b, g)) -> RandomM g b
forall a b. (a -> b) -> a -> b
$ \g
g -> let (a
x, g
g') = RandomM g a -> g -> (a, g)
forall g a. RandomM g a -> g -> (a, g)
runRandomM RandomM g a
r g
g
                               in (a -> b
f a
x, g
g')

instance Applicative (RandomM g) where
    pure :: forall a. a -> RandomM g a
pure a
x = (g -> (a, g)) -> RandomM g a
forall g a. (g -> (a, g)) -> RandomM g a
RandomM (a
x,)
    RandomM g (a -> b)
rf <*> :: forall a b. RandomM g (a -> b) -> RandomM g a -> RandomM g b
<*> RandomM g a
r = (g -> (b, g)) -> RandomM g b
forall g a. (g -> (a, g)) -> RandomM g a
RandomM ((g -> (b, g)) -> RandomM g b) -> (g -> (b, g)) -> RandomM g b
forall a b. (a -> b) -> a -> b
$ \g
g -> let (a
x, g
g') = RandomM g a -> g -> (a, g)
forall g a. RandomM g a -> g -> (a, g)
runRandomM RandomM g a
r g
g
                                   (a -> b
f, g
g'') = RandomM g (a -> b) -> g -> (a -> b, g)
forall g a. RandomM g a -> g -> (a, g)
runRandomM RandomM g (a -> b)
rf g
g'
                               in (a -> b
f a
x, g
g'')

instance Monad (RandomM g) where
    RandomM g a
r >>= :: forall a b. RandomM g a -> (a -> RandomM g b) -> RandomM g b
>>= a -> RandomM g b
f = (g -> (b, g)) -> RandomM g b
forall g a. (g -> (a, g)) -> RandomM g a
RandomM ((g -> (b, g)) -> RandomM g b) -> (g -> (b, g)) -> RandomM g b
forall a b. (a -> b) -> a -> b
$ \g
g -> let (a
x, g
g') = RandomM g a -> g -> (a, g)
forall g a. RandomM g a -> g -> (a, g)
runRandomM RandomM g a
r g
g
                              in RandomM g b -> g -> (b, g)
forall g a. RandomM g a -> g -> (a, g)
runRandomM (a -> RandomM g b
f a
x) g
g'

-- | Runs a random monad using the global RNG.
runRandomMIO :: RandomM StdGen a -> IO a
runRandomMIO :: forall a. RandomM StdGen a -> IO a
runRandomMIO RandomM StdGen a
r = (a, StdGen) -> a
forall a b. (a, b) -> a
fst ((a, StdGen) -> a) -> (StdGen -> (a, StdGen)) -> StdGen -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RandomM StdGen a -> StdGen -> (a, StdGen)
forall g a. RandomM g a -> g -> (a, g)
runRandomM RandomM StdGen a
r (StdGen -> a) -> IO StdGen -> IO a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO StdGen
forall (m :: * -> *). MonadIO m => m StdGen
getStdGen

-- | Fetches the generator inside the random monad.
getRandomGen :: RandomM g g
getRandomGen :: forall g. RandomM g g
getRandomGen = (g -> (g, g)) -> RandomM g g
forall g a. (g -> (a, g)) -> RandomM g a
RandomM ((g -> (g, g)) -> RandomM g g) -> (g -> (g, g)) -> RandomM g g
forall a b. (a -> b) -> a -> b
$ \g
g -> (g
g, g
g)

-- | Generates a random value inside the random monad.
randomM :: (RandomGen g, Random a) => RandomM g a
randomM :: forall g a. (RandomGen g, Random a) => RandomM g a
randomM = (g -> (a, g)) -> RandomM g a
forall g a. (g -> (a, g)) -> RandomM g a
RandomM g -> (a, g)
forall a g. (Random a, RandomGen g) => g -> (a, g)
forall g. RandomGen g => g -> (a, g)
random

-- | Generates a bounded random value inside the random monad.
randomRM :: (RandomGen g, Random a) => (a, a) -> RandomM g a
randomRM :: forall g a. (RandomGen g, Random a) => (a, a) -> RandomM g a
randomRM (a, a)
r = (g -> (a, g)) -> RandomM g a
forall g a. (g -> (a, g)) -> RandomM g a
RandomM ((g -> (a, g)) -> RandomM g a) -> (g -> (a, g)) -> RandomM g a
forall a b. (a -> b) -> a -> b
$ (a, a) -> g -> (a, g)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
forall g. RandomGen g => (a, a) -> g -> (a, g)
randomR (a, a)
r

-- | Generates n random values in the given range without consuming the generator.
nRandomsR :: (RandomGen g, Random a) => Int -> ([a], [a]) -> RandomM g [a]
nRandomsR :: forall g a.
(RandomGen g, Random a) =>
Int -> ([a], [a]) -> RandomM g [a]
nRandomsR Int
0 ([a], [a])
_ = [a] -> RandomM g [a]
forall a. a -> RandomM g a
forall (m :: * -> *) a. Monad m => a -> m a
return []
nRandomsR Int
n (a
l:[a]
ls, a
h:[a]
hs) = do
    a
x <- (a, a) -> RandomM g a
forall g a. (RandomGen g, Random a) => (a, a) -> RandomM g a
randomRM (a
l, a
h)
    (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:) ([a] -> [a]) -> RandomM g [a] -> RandomM g [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> ([a], [a]) -> RandomM g [a]
forall g a.
(RandomGen g, Random a) =>
Int -> ([a], [a]) -> RandomM g [a]
nRandomsR (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ([a]
ls, [a]
hs)
nRandomsR Int
n ([a], [a])
_ = [Char] -> RandomM g [a]
forall a. HasCallStack => [Char] -> a
error ([Char] -> RandomM g [a]) -> [Char] -> RandomM g [a]
forall a b. (a -> b) -> a -> b
$ [Char]
"No range(s) for remaining " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
n [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" random value(s)!"

-- | Generates n random values without consuming the generator.
nRandoms :: (RandomGen g, Random a) => Int -> RandomM g [a]
nRandoms :: forall g a. (RandomGen g, Random a) => Int -> RandomM g [a]
nRandoms Int
0 = [a] -> RandomM g [a]
forall a. a -> RandomM g a
forall (m :: * -> *) a. Monad m => a -> m a
return []
nRandoms Int
n = do
    a
x <- RandomM g a
forall g a. (RandomGen g, Random a) => RandomM g a
randomM
    (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:) ([a] -> [a]) -> RandomM g [a] -> RandomM g [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> RandomM g [a]
forall g a. (RandomGen g, Random a) => Int -> RandomM g [a]
nRandoms (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)