-- | Representation of probabilities and random computations.
module Game.LambdaHack.Core.Random
  ( -- * The @Rng@ monad
    Rnd
    -- * Random operations
  , randomR, random, oneOf, shuffle, frequency
    -- * Fractional chance
  , Chance, chance
    -- * Casting dice scaled with level
  , castDice, oddsDice, castDiceXY
    -- * Specialized monadic folds
  , foldrM, foldlM'
#ifdef EXPOSE_INTERNAL
    -- * Internal operations
  , rollFreq
#endif
  ) where

import Prelude ()

import Game.LambdaHack.Core.Prelude

import qualified Control.Monad.Trans.State.Strict as St
import           Data.Ratio
import qualified System.Random as R

import qualified Game.LambdaHack.Core.Dice as Dice
import           Game.LambdaHack.Core.Frequency

-- | The monad of computations with random generator state.
type Rnd a = St.State R.StdGen a

-- | Get a random object within a range with a uniform distribution.
randomR :: (R.Random a) => (a, a) -> Rnd a
{-# INLINE randomR #-}
randomR :: (a, a) -> Rnd a
randomR = (StdGen -> (a, StdGen)) -> Rnd a
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
St.state ((StdGen -> (a, StdGen)) -> Rnd a)
-> ((a, a) -> StdGen -> (a, StdGen)) -> (a, a) -> Rnd a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, a) -> StdGen -> (a, StdGen)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
R.randomR

-- | Get a random object of a given type with a uniform distribution.
random :: (R.Random a) => Rnd a
{-# INLINE random #-}
random :: Rnd a
random = (StdGen -> (a, StdGen)) -> Rnd a
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
St.state StdGen -> (a, StdGen)
forall a g. (Random a, RandomGen g) => g -> (a, g)
R.random

-- | Get any element of a list with equal probability.
oneOf :: [a] -> Rnd a
oneOf :: [a] -> Rnd a
oneOf [] = [Char] -> Rnd a
forall a. HasCallStack => [Char] -> a
error ([Char] -> Rnd a) -> [Char] -> Rnd a
forall a b. (a -> b) -> a -> b
$ "oneOf []" [Char] -> () -> [Char]
forall v. Show v => [Char] -> v -> [Char]
`showFailure` ()
oneOf [x :: a
x] = a -> Rnd a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x
oneOf xs :: [a]
xs = do
  Int
r <- (Int, Int) -> Rnd Int
forall a. Random a => (a, a) -> Rnd a
randomR (0, [a] -> Int
forall a. [a] -> Int
length [a]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1)
  a -> Rnd a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Rnd a) -> a -> Rnd a
forall a b. (a -> b) -> a -> b
$! [a]
xs [a] -> Int -> a
forall a. [a] -> Int -> a
!! Int
r

-- | Generates a random permutation. Naive, but good enough for small inputs.
shuffle :: Eq a => [a] -> Rnd [a]
shuffle :: [a] -> Rnd [a]
shuffle [] = [a] -> Rnd [a]
forall (m :: * -> *) a. Monad m => a -> m a
return []
shuffle l :: [a]
l = do
  a
x <- [a] -> Rnd a
forall a. [a] -> Rnd a
oneOf [a]
l
  (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
:) ([a] -> [a]) -> Rnd [a] -> Rnd [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [a] -> Rnd [a]
forall a. Eq a => [a] -> Rnd [a]
shuffle (a -> [a] -> [a]
forall a. Eq a => a -> [a] -> [a]
delete a
x [a]
l)

-- | Gen an element according to a frequency distribution.
frequency :: Show a => Frequency a -> Rnd a
{-# INLINE frequency #-}
frequency :: Frequency a -> Rnd a
frequency = (StdGen -> (a, StdGen)) -> Rnd a
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
St.state ((StdGen -> (a, StdGen)) -> Rnd a)
-> (Frequency a -> StdGen -> (a, StdGen)) -> Frequency a -> Rnd a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Frequency a -> StdGen -> (a, StdGen)
forall a. Show a => Frequency a -> StdGen -> (a, StdGen)
rollFreq

-- | Randomly choose an item according to the distribution.
rollFreq :: Show a => Frequency a -> R.StdGen -> (a, R.StdGen)
rollFreq :: Frequency a -> StdGen -> (a, StdGen)
rollFreq fr :: Frequency a
fr g :: StdGen
g = case Frequency a -> [(Int, a)]
forall a. Frequency a -> [(Int, a)]
runFrequency Frequency a
fr of
  [] -> [Char] -> (a, StdGen)
forall a. HasCallStack => [Char] -> a
error ([Char] -> (a, StdGen)) -> [Char] -> (a, StdGen)
forall a b. (a -> b) -> a -> b
$ "choice from an empty frequency"
                [Char] -> Text -> [Char]
forall v. Show v => [Char] -> v -> [Char]
`showFailure` Frequency a -> Text
forall a. Frequency a -> Text
nameFrequency Frequency a
fr
  [(n :: Int
n, x :: a
x)] | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= 0 -> [Char] -> (a, StdGen)
forall a. HasCallStack => [Char] -> a
error ([Char] -> (a, StdGen)) -> [Char] -> (a, StdGen)
forall a b. (a -> b) -> a -> b
$ "singleton void frequency"
                               [Char] -> (Text, Int, a) -> [Char]
forall v. Show v => [Char] -> v -> [Char]
`showFailure` (Frequency a -> Text
forall a. Frequency a -> Text
nameFrequency Frequency a
fr, Int
n, a
x)
  [(_, x :: a
x)] -> (a
x, StdGen
g)  -- speedup
  fs :: [(Int, a)]
fs -> let sumf :: Int
sumf = (Int -> (Int, a) -> Int) -> Int -> [(Int, a)] -> Int
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\ !Int
acc (!Int
n, _) -> Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) 0 [(Int, a)]
fs
            (r :: Int
r, ng :: StdGen
ng) = (Int, Int) -> StdGen -> (Int, StdGen)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
R.randomR (1, Int
sumf) StdGen
g
            frec :: Int -> [(Int, a)] -> a
            frec :: Int -> [(Int, a)] -> a
frec !Int
m [] = [Char] -> a
forall a. HasCallStack => [Char] -> a
error ([Char] -> a) -> [Char] -> a
forall a b. (a -> b) -> a -> b
$ "impossible roll"
                                 [Char] -> (Text, [(Int, a)], Int) -> [Char]
forall v. Show v => [Char] -> v -> [Char]
`showFailure` (Frequency a -> Text
forall a. Frequency a -> Text
nameFrequency Frequency a
fr, [(Int, a)]
fs, Int
m)
            frec m :: Int
m ((n :: Int
n, x :: a
x) : _) | Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n = a
x
            frec m :: Int
m ((n :: Int
n, _) : xs :: [(Int, a)]
xs) = Int -> [(Int, a)] -> a
forall a. Int -> [(Int, a)] -> a
frec (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n) [(Int, a)]
xs
        in Bool -> (a, StdGen) -> (a, StdGen)
forall a. HasCallStack => Bool -> a -> a
assert (Int
sumf Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> 0 Bool -> ([Char], (Text, [(Int, a)])) -> Bool
forall a. Show a => Bool -> a -> Bool
`blame` "frequency with nothing to pick"
                            [Char] -> (Text, [(Int, a)]) -> ([Char], (Text, [(Int, a)]))
forall v. [Char] -> v -> ([Char], v)
`swith` (Frequency a -> Text
forall a. Frequency a -> Text
nameFrequency Frequency a
fr, [(Int, a)]
fs))
             (Int -> [(Int, a)] -> a
forall a. Int -> [(Int, a)] -> a
frec Int
r [(Int, a)]
fs, StdGen
ng)

-- | Fractional chance.
type Chance = Rational

-- | Give @True@, with probability determined by the fraction.
chance :: Chance -> Rnd Bool
chance :: Chance -> Rnd Bool
chance r :: Chance
r = do
  let n :: Integer
n = Chance -> Integer
forall a. Ratio a -> a
numerator Chance
r
      d :: Integer
d = Chance -> Integer
forall a. Ratio a -> a
denominator Chance
r
  Integer
k <- (Integer, Integer) -> Rnd Integer
forall a. Random a => (a, a) -> Rnd a
randomR (1, Integer
d)
  Bool -> Rnd Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Integer
k Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
n)

-- | Cast dice scaled with current level depth.
-- Note that at the first level, the scaled dice are always ignored.
castDice :: Dice.AbsDepth -> Dice.AbsDepth -> Dice.Dice -> Rnd Int
castDice :: AbsDepth -> AbsDepth -> Dice -> Rnd Int
castDice = ((Int, Int) -> Rnd Int) -> AbsDepth -> AbsDepth -> Dice -> Rnd Int
forall (m :: * -> *).
Monad m =>
((Int, Int) -> m Int) -> AbsDepth -> AbsDepth -> Dice -> m Int
Dice.castDice (Int, Int) -> Rnd Int
forall a. Random a => (a, a) -> Rnd a
randomR

-- | Cast dice scaled with current level depth and return @True@
-- if the results is greater than 50.
oddsDice :: Dice.AbsDepth -> Dice.AbsDepth -> Dice.Dice -> Rnd Bool
oddsDice :: AbsDepth -> AbsDepth -> Dice -> Rnd Bool
oddsDice ldepth :: AbsDepth
ldepth totalDepth :: AbsDepth
totalDepth dice :: Dice
dice = do
  Int
c <- AbsDepth -> AbsDepth -> Dice -> Rnd Int
castDice AbsDepth
ldepth AbsDepth
totalDepth Dice
dice
  Bool -> Rnd Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> Rnd Bool) -> Bool -> Rnd Bool
forall a b. (a -> b) -> a -> b
$! Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> 50

-- | Cast dice, scaled with current level depth, for coordinates.
castDiceXY :: Dice.AbsDepth -> Dice.AbsDepth -> Dice.DiceXY -> Rnd (Int, Int)
castDiceXY :: AbsDepth -> AbsDepth -> DiceXY -> Rnd (Int, Int)
castDiceXY ldepth :: AbsDepth
ldepth totalDepth :: AbsDepth
totalDepth (Dice.DiceXY dx :: Dice
dx dy :: Dice
dy) = do
  Int
x <- AbsDepth -> AbsDepth -> Dice -> Rnd Int
castDice AbsDepth
ldepth AbsDepth
totalDepth Dice
dx
  Int
y <- AbsDepth -> AbsDepth -> Dice -> Rnd Int
castDice AbsDepth
ldepth AbsDepth
totalDepth Dice
dy
  (Int, Int) -> Rnd (Int, Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
x, Int
y)

foldrM :: Foldable t => (a -> b -> Rnd b) -> b -> t a -> Rnd b
foldrM :: (a -> b -> Rnd b) -> b -> t a -> Rnd b
foldrM f :: a -> b -> Rnd b
f z0 :: b
z0 xs :: t a
xs = let f' :: a -> (b, StdGen) -> (b, StdGen)
f' x :: a
x (z :: b
z, g :: StdGen
g) = Rnd b -> StdGen -> (b, StdGen)
forall s a. State s a -> s -> (a, s)
St.runState (a -> b -> Rnd b
f a
x b
z) StdGen
g
                 in (StdGen -> (b, StdGen)) -> Rnd b
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
St.state ((StdGen -> (b, StdGen)) -> Rnd b)
-> (StdGen -> (b, StdGen)) -> Rnd b
forall a b. (a -> b) -> a -> b
$ \g :: StdGen
g -> (a -> (b, StdGen) -> (b, StdGen))
-> (b, StdGen) -> t a -> (b, StdGen)
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> (b, StdGen) -> (b, StdGen)
f' (b
z0, StdGen
g) t a
xs

foldlM' :: Foldable t => (b -> a -> Rnd b) -> b -> t a -> Rnd b
foldlM' :: (b -> a -> Rnd b) -> b -> t a -> Rnd b
foldlM' f :: b -> a -> Rnd b
f z0 :: b
z0 xs :: t a
xs = let f' :: (b, StdGen) -> a -> (b, StdGen)
f' (z :: b
z, g :: StdGen
g) x :: a
x = Rnd b -> StdGen -> (b, StdGen)
forall s a. State s a -> s -> (a, s)
St.runState (b -> a -> Rnd b
f b
z a
x) StdGen
g
                  in (StdGen -> (b, StdGen)) -> Rnd b
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
St.state ((StdGen -> (b, StdGen)) -> Rnd b)
-> (StdGen -> (b, StdGen)) -> Rnd b
forall a b. (a -> b) -> a -> b
$ \g :: StdGen
g -> ((b, StdGen) -> a -> (b, StdGen))
-> (b, StdGen) -> t a -> (b, StdGen)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (b, StdGen) -> a -> (b, StdGen)
f' (b
z0, StdGen
g) t a
xs