-- | Randomized operations using a `Random` monad for "PCG" operations.
module MCSP.System.Random (
    -- * Random Monad
    Random,
    evalRandom,
    liftRandom,
    lazyRandom,
    generate,
    generateFast,

    -- * Evaluation
    Seed,
    readSeedP,
    readSeed,
    showSeedS,
    showSeed,
    generateFastWith,
    generateWith,
    randomSeed,

    -- * Random Values
    PCG.Variate,
    uniform,
    uniformR,
    uniformB,
    uniformE,
    uniformRE,
    choice,
    weightedChoice,
    shuffle,
    partitions,
    repeatR,
    iterateR,
) where

import Control.Applicative (pure)
import Control.Monad (mapM)
import Control.Monad.ST (runST)
import Data.Foldable (length, sum)
import Data.Function (($), (.))
import Data.Int (Int)
import Data.List (map)
import Data.List.NonEmpty (NonEmpty (..), head, nonEmpty, (<|))
import Data.List.NonEmpty.Extra ((!?))
import Data.Maybe (Maybe (..), fromMaybe)
import Data.Tuple (fst, snd)
import Data.Vector (Vector, (!))
import Data.Vector.Algorithms.Search (binarySearchL)
import Data.Vector.Generic qualified as Vector (Vector, init, length, splitAt, unsafeThaw)
import Data.Vector.Unboxed qualified as Unboxed (Vector)
import Data.Word (Word32)
import GHC.Enum (Bounded (..), Enum (..))
import GHC.Exts (IsList (..))
import GHC.Float (Double)
import GHC.Num ((*), (+), (-))
import GHC.Real (fromIntegral, truncate, (/))
import System.Random.PCG qualified as PCG (Variate (..))
import System.Random.Shuffle qualified as Shuffle (shuffle)

import MCSP.System.Random.Generate (
    Seed,
    generate,
    generateFast,
    generateFastWith,
    generateWith,
    randomSeed,
    readSeed,
    readSeedP,
    showSeed,
    showSeedS,
 )
import MCSP.System.Random.Monad (
    Random,
    evalRandom,
    lazyRandom,
    liftRandom,
 )

-- ------------- --
-- Random Values --
-- ------------- --

-- | /O(1)/ Generate a uniformly distributed random variate.
--
-- * Use entire range for integral types.
-- * Use (0,1] range for floating types.
--
-- >>> import Prelude (Double)
-- >>> generateWith @Double (1,2) uniform
-- 0.6502342391751404
uniform :: PCG.Variate a => Random a
uniform :: forall a. Variate a => Random a
uniform = (forall g (m :: * -> *). Generator g m => g -> m a) -> Random a
forall a.
(forall g (m :: * -> *). Generator g m => g -> m a) -> Random a
liftRandom g -> m a
forall a g (m :: * -> *). (Variate a, Generator g m) => g -> m a
forall g (m :: * -> *). Generator g m => g -> m a
PCG.uniform
{-# INLINE uniform #-}

-- | /O(1)/ Generate a uniformly distributed random variate in the given range.
--
-- * Use inclusive range for integral types.
-- * Use (a,b] range for floating types.
--
-- >>> generateWith (1,2) $ uniformR @Int 10 50
-- 16
uniformR :: PCG.Variate a => a -> a -> Random a
uniformR :: forall a. Variate a => a -> a -> Random a
uniformR a
lo a
hi = (forall g (m :: * -> *). Generator g m => g -> m a) -> Random a
forall a.
(forall g (m :: * -> *). Generator g m => g -> m a) -> Random a
liftRandom ((forall g (m :: * -> *). Generator g m => g -> m a) -> Random a)
-> (forall g (m :: * -> *). Generator g m => g -> m a) -> Random a
forall a b. (a -> b) -> a -> b
$ (a, a) -> g -> m a
forall a g (m :: * -> *).
(Variate a, Generator g m) =>
(a, a) -> g -> m a
forall g (m :: * -> *). Generator g m => (a, a) -> g -> m a
PCG.uniformR (a
lo, a
hi)
{-# INLINE uniformR #-}

-- | /O(1)/ Generate a uniformly distributed random variate in the range [0,b).
--
-- * For integral types the bound must be less than the max bound of `Data.Word.Word32`
-- (4294967295). Behaviour is undefined for negative bounds.
--
-- >>> generateWith (1,2) $ uniformB @Int 200
-- 143
uniformB :: PCG.Variate a => a -> Random a
uniformB :: forall a. Variate a => a -> Random a
uniformB a
b = (forall g (m :: * -> *). Generator g m => g -> m a) -> Random a
forall a.
(forall g (m :: * -> *). Generator g m => g -> m a) -> Random a
liftRandom ((forall g (m :: * -> *). Generator g m => g -> m a) -> Random a)
-> (forall g (m :: * -> *). Generator g m => g -> m a) -> Random a
forall a b. (a -> b) -> a -> b
$ a -> g -> m a
forall a g (m :: * -> *).
(Variate a, Generator g m) =>
a -> g -> m a
forall g (m :: * -> *). Generator g m => a -> g -> m a
PCG.uniformB a
b
{-# INLINE uniformB #-}

-- | /O(1)/ Generate a uniformly distributed random variate.
--
-- It should be equivalent to @`uniformR` (`minBound`, `maxBound`)@, but for non-`PCG.Variate`.
--
-- >>> import Prelude (Show)
-- >>> data T = A | B | C | D deriving (Enum, Bounded, Show)
-- >>> generateWith @T (1,3) uniformE
-- C
uniformE :: (Enum a, Bounded a) => Random a
uniformE :: forall a. (Enum a, Bounded a) => Random a
uniformE = a -> a -> Random a
forall a. Enum a => a -> a -> Random a
uniformRE a
forall a. Bounded a => a
minBound a
forall a. Bounded a => a
maxBound
{-# INLINE uniformE #-}

-- | /O(1)/ Generate a uniformly distributed random variate in the given inclusive range.
--
-- It should be equivalent to `uniformR`, but for non-`PCG.Variate`.
--
-- >>> generateWith (1,3) $ uniformRE @Int 0 10
-- 5
--
-- >>> generateWith (1,3) $ uniformRE 'a' 'z'
-- 'n'
uniformRE :: Enum a => a -> a -> Random a
uniformRE :: forall a. Enum a => a -> a -> Random a
uniformRE a
lo a
hi = do
    let loNum :: Int
loNum = a -> Int
forall a. Enum a => a -> Int
fromEnum a
lo
    let hiNum :: Int
hiNum = a -> Int
forall a. Enum a => a -> Int
fromEnum a
hi
    Int
value <- Int -> Int -> Random Int
forall a. Variate a => a -> a -> Random a
uniformR Int
loNum Int
hiNum
    pure (Int -> a
forall a. Enum a => Int -> a
toEnum Int
value)
{-# INLINE uniformRE #-}

-- | /O(1)/ Choose a single random value from a non-empty list.
--
-- >>> generateWith (1,2) $ choice ["hi", "hello", "ola"]
-- "hello"
choice :: NonEmpty a -> Random a
choice :: forall a. NonEmpty a -> Random a
choice NonEmpty a
values = do
    Int
i <- Int -> Random Int
forall a. Variate a => a -> Random a
uniformB (NonEmpty a -> Int
forall a. NonEmpty a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length NonEmpty a
values)
    pure $ a -> Maybe a -> a
forall a. a -> Maybe a -> a
fromMaybe (NonEmpty a -> a
forall {a}. NonEmpty a -> a
headNE NonEmpty a
values) (NonEmpty a
values NonEmpty a -> Int -> Maybe a
forall a. NonEmpty a -> Int -> Maybe a
!? Int
i)
  where
    headNE :: NonEmpty a -> a
headNE = NonEmpty a -> a
forall {a}. NonEmpty a -> a
head
{-# INLINE choice #-}

-- | /O(n)/ Generate weighted access table for a list of values.
--
-- >>> tabulate [(1, "hi"), (10, "hello")]
-- ([390451572],["hi","hello"])
tabulate :: NonEmpty (Double, a) -> (Unboxed.Vector Word32, Vector a)
tabulate :: forall a. NonEmpty (Double, a) -> (Vector Word32, Vector a)
tabulate (NonEmpty (Double, a) -> [Item (NonEmpty (Double, a))]
forall l. IsList l => l -> [Item l]
toList -> [Item (NonEmpty (Double, a))]
values) =
    ( Vector Word32 -> Vector Word32
forall (v :: * -> *) a. Vector v a => v a -> v a
Vector.init (((Double, a) -> Item (Vector Word32))
-> [(Double, a)] -> Vector Word32
forall {c} {a}. IsList c => (a -> Item c) -> [a] -> c
vectorOf (Double -> Word32
toW32 (Double -> Word32)
-> ((Double, a) -> Double) -> (Double, a) -> Word32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double, a) -> Double
forall a b. (a, b) -> a
fst) [(Double, a)]
[Item (NonEmpty (Double, a))]
values),
      ((Double, a) -> Item (Vector a)) -> [(Double, a)] -> Vector a
forall {c} {a}. IsList c => (a -> Item c) -> [a] -> c
vectorOf (Double, a) -> a
(Double, a) -> Item (Vector a)
forall a b. (a, b) -> b
snd [(Double, a)]
[Item (NonEmpty (Double, a))]
values
    )
  where
    vectorOf :: (a -> Item c) -> [a] -> c
vectorOf a -> Item c
f = [Item c] -> c
forall l. IsList l => [Item l] -> l
fromList ([Item c] -> c) -> ([a] -> [Item c]) -> [a] -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Item c) -> [a] -> [Item c]
forall a b. (a -> b) -> [a] -> [b]
map a -> Item c
f
    maxWeight :: Double
maxWeight = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (((Double, a) -> Double) -> [(Double, a)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double, a) -> Double
forall a b. (a, b) -> a
fst [(Double, a)]
[Item (NonEmpty (Double, a))]
values)
    maxW32 :: Double
maxW32 = forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word32 Word32
forall a. Bounded a => a
maxBound
    toW32 :: Double -> Word32
toW32 Double
x = Double -> Word32
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
truncate (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
maxW32 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
maxWeight)
{-# INLINE tabulate #-}

-- | /O(n)/ Choose a single random value from a non-empty with probability proportional to weight.
--
-- >>> import Control.Monad (replicateM)
-- >>> generateWith (1,2) $ replicateM 10 $ weightedChoice [(1, "hi"), (10, "hello")]
-- ["hello","hello","hello","hello","hi","hello","hello","hello","hello","hello"]
weightedChoice :: NonEmpty (Double, a) -> Random a
weightedChoice :: forall a. NonEmpty (Double, a) -> Random a
weightedChoice (NonEmpty (Double, a) -> (Vector Word32, Vector a)
forall a. NonEmpty (Double, a) -> (Vector Word32, Vector a)
tabulate -> (Vector Word32
positions, Vector a
values)) = do
    Word32
weight <- Random Word32
forall a. Variate a => Random a
uniform
    let idx :: Int
idx = Vector Word32 -> Word32 -> Int
forall {v :: * -> *} {a}. (Vector v a, Ord a) => v a -> a -> Int
binarySearch Vector Word32
positions Word32
weight
    a -> Random a
forall a. a -> Random a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vector a
values Vector a -> Int -> a
forall a. Vector a -> Int -> a
! Int
idx)
  where
    binarySearch :: v a -> a -> Int
binarySearch v a
v a
x = (forall s. ST s Int) -> Int
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s Int) -> Int) -> (forall s. ST s Int) -> Int
forall a b. (a -> b) -> a -> b
$ do
        -- SAFETY: binarySearchL does NOT modify the vector,
        -- I don't know why they chose to expose a mutable API only
        Mutable v s a
mv <- v a -> ST s (Mutable v (PrimState (ST s)) a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
v a -> m (Mutable v (PrimState m) a)
Vector.unsafeThaw v a
v
        Mutable v (PrimState (ST s)) a -> a -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e, Ord e) =>
v (PrimState m) e -> e -> m Int
binarySearchL Mutable v s a
Mutable v (PrimState (ST s)) a
mv a
x
{-# INLINE weightedChoice #-}

-- | /O(n)/ Generates indices for `S.shuffle`.
--
-- See [random-shuffle](https://hackage.haskell.org/package/random-shuffle-0.0.4/docs/System-Random-Shuffle.html#v:shuffle).
treeIndices :: Int -> Random [Int]
treeIndices :: Int -> Random [Int]
treeIndices Int
n = (Int -> Random Int) -> [Int] -> Random [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Int -> Random Int
forall a. Variate a => a -> Random a
uniformB [Int
Item [Int]
n, Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 .. Int
Item [Int]
2]
{-# INLINE treeIndices #-}

-- | /O(?)/ Shuffles a non-empty list.
--
-- This simple implementation runs forever with empty lists.
shuffleNonEmpty :: IsList l => NonEmpty (Item l) -> Random l
shuffleNonEmpty :: forall l. IsList l => NonEmpty (Item l) -> Random l
shuffleNonEmpty (Item l
h :| [Item l]
rest) = do
    let xs :: [Item l]
xs = Item l
h Item l -> [Item l] -> [Item l]
forall a. a -> [a] -> [a]
: [Item l]
rest
    let n :: Int
n = [Item l] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Item l]
xs
    [Int]
idx <- Int -> Random [Int]
treeIndices Int
n
    let ys :: [Item l]
ys = [Item l] -> [Int] -> [Item l]
forall a. [a] -> [Int] -> [a]
Shuffle.shuffle [Item l]
xs [Int]
idx
    l -> Random l
forall a. a -> Random a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> [Item l] -> l
forall l. IsList l => Int -> [Item l] -> l
fromListN Int
n [Item l]
ys)
{-# INLINE shuffleNonEmpty #-}

-- | /O(?)/ Shuffles a list randomly.
--
-- >>> generateWith (1,2) $ shuffle @[Int] [1..5]
-- [4,1,2,3,5]
shuffle :: IsList l => l -> Random l
shuffle :: forall l. IsList l => l -> Random l
shuffle l
values = case [Item l] -> Maybe (NonEmpty (Item l))
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty (l -> [Item l]
forall l. IsList l => l -> [Item l]
toList l
values) of
    Just NonEmpty (Item l)
xs -> NonEmpty (Item l) -> Random l
forall l. IsList l => NonEmpty (Item l) -> Random l
shuffleNonEmpty NonEmpty (Item l)
xs
    Maybe (NonEmpty (Item l))
Nothing -> l -> Random l
forall a. a -> Random a
forall (f :: * -> *) a. Applicative f => a -> f a
pure l
values
{-# INLINEABLE shuffle #-}

-- | /O(n)/ Generate random partitions of a vector.
--
-- >>> generateWith (1,4) $ partitions @Vector @Int [1..10]
-- [[1,2,3],[4,5,6,7,8],[9],[10]]
partitions :: Vector.Vector v a => v a -> Random [v a]
partitions :: forall (v :: * -> *) a. Vector v a => v a -> Random [v a]
partitions v a
xs = case v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
Vector.length v a
xs of
    Int
0 -> [v a] -> Random [v a]
forall a. a -> Random a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    Int
1 -> [v a] -> Random [v a]
forall a. a -> Random a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [v a
Item [v a]
xs]
    Int
n -> do
        Int
idx <- Int -> Random Int
forall a. Variate a => a -> Random a
uniformB Int
n
        let (v a
part, v a
rest) = Int -> v a -> (v a, v a)
forall (v :: * -> *) a. Vector v a => Int -> v a -> (v a, v a)
Vector.splitAt (Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) v a
xs
        [v a]
parts <- v a -> Random [v a]
forall (v :: * -> *) a. Vector v a => v a -> Random [v a]
partitions v a
rest
        pure (v a
part v a -> [v a] -> [v a]
forall a. a -> [a] -> [a]
: [v a]
parts)
{-# INLINEABLE partitions #-}

-- | Generates an infinite list of random values.
--
-- >>> import Prelude (take, (<$>))
-- >>> generateWith (1,3) (take 3 <$> repeatR (uniform @Int))
-- [-8858759364290496978,-2446124676014956441,1387719708863309077]
repeatR :: Random a -> Random [a]
repeatR :: forall a. Random a -> Random [a]
repeatR Random a
r = Random [a] -> Random [a]
forall a. Random a -> Random a
lazyRandom (Random [a] -> Random [a]) -> Random [a] -> Random [a]
forall a b. (a -> b) -> a -> b
$ do
    a
value <- Random a
r
    [a]
rest <- Random a -> Random [a]
forall a. Random a -> Random [a]
repeatR Random a
r
    pure (a
value a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
rest)
{-# INLINEABLE repeatR #-}

-- | Random version of `Data.List.iterate`.
--
-- >>> import Prelude ((<$>), (.), liftA2)
-- >>> import Data.List.NonEmpty (take)
-- >>> generateWith (1,3) $ take 3 <$> iterateR (liftA2 (+) uniform . pure) 2
-- [2.0,2.0197656927151786,2.3871610084947057]
iterateR :: (a -> Random a) -> a -> Random (NonEmpty a)
iterateR :: forall a. (a -> Random a) -> a -> Random (NonEmpty a)
iterateR a -> Random a
next a
value = Random (NonEmpty a) -> Random (NonEmpty a)
forall a. Random a -> Random a
lazyRandom (Random (NonEmpty a) -> Random (NonEmpty a))
-> Random (NonEmpty a) -> Random (NonEmpty a)
forall a b. (a -> b) -> a -> b
$ do
    a
newValue <- a -> Random a
next a
value
    NonEmpty a
rest <- (a -> Random a) -> a -> Random (NonEmpty a)
forall a. (a -> Random a) -> a -> Random (NonEmpty a)
iterateR a -> Random a
next a
newValue
    pure (a
value a -> NonEmpty a -> NonEmpty a
forall a. a -> NonEmpty a -> NonEmpty a
<| NonEmpty a
rest)
{-# INLINEABLE iterateR #-}