-- | Representation of probabilities and random computations.
module Game.LambdaHack.Core.Random
  ( -- * The @Rng@ monad
    Rnd
    -- * Random operations
  , randomR, randomR0, nextRandom, randomWord32
  , oneOf, shuffle, invalidInformationCode, shuffleExcept, 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.Int (Int32)
import           Data.Ratio
import qualified Data.Vector.Unboxed as U
import           Data.Word (Word16, Word32)
import qualified System.Random.SplitMix32 as SM

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 SM.SMGen a

-- | Get a random object within a (inclusive) range with a uniform distribution.
randomR :: (Integral a) => (a, a) -> Rnd a
{-# INLINE randomR #-}
randomR :: (a, a) -> Rnd a
randomR (a
0, a
h) = a -> Rnd a
forall a. Integral a => a -> Rnd a
randomR0 a
h
randomR (a
l, a
h) | a
l a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
h = [Char] -> Rnd a
forall a. HasCallStack => [Char] -> a
error [Char]
"randomR: empty range"
randomR (a
l, a
h) = (SMGen -> (a, SMGen)) -> Rnd a
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
St.state ((SMGen -> (a, SMGen)) -> Rnd a) -> (SMGen -> (a, SMGen)) -> Rnd a
forall a b. (a -> b) -> a -> b
$ \SMGen
g ->
  let (a
x, SMGen
g') = a -> SMGen -> (a, SMGen)
forall a. Integral a => a -> SMGen -> (a, SMGen)
nextRandom (a
h a -> a -> a
forall a. Num a => a -> a -> a
- a
l) SMGen
g
  in (a
x a -> a -> a
forall a. Num a => a -> a -> a
+ a
l, SMGen
g')

-- | Generate random 'Integral' in @[0, x]@ range.
randomR0 :: (Integral a) => a -> Rnd a
{-# INLINE randomR0 #-}
randomR0 :: a -> Rnd a
randomR0 a
h = (SMGen -> (a, SMGen)) -> Rnd a
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
St.state ((SMGen -> (a, SMGen)) -> Rnd a) -> (SMGen -> (a, SMGen)) -> Rnd a
forall a b. (a -> b) -> a -> b
$ a -> SMGen -> (a, SMGen)
forall a. Integral a => a -> SMGen -> (a, SMGen)
nextRandom a
h

-- | Generate a random integral value in @[0, x]@ range, where @x@ is within
-- @Int32@.
--
-- The limitation to @Int32@ values is needed to keep it working on signed
-- types. In package @random@, a much more complex scheme is used
-- to keep it working for arbitrary fixed number of bits.
nextRandom :: forall a. (Integral a) => a -> SM.SMGen -> (a, SM.SMGen)
{-# INLINE nextRandom #-}
nextRandom :: a -> SMGen -> (a, SMGen)
nextRandom a
0 SMGen
g = (a
0, SMGen
g)
nextRandom a
h SMGen
g = Bool -> (a, SMGen) -> (a, SMGen)
forall a. HasCallStack => Bool -> a -> a
assert (a
h a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
0 Bool -> Bool -> Bool
&& a -> Integer
forall a. Integral a => a -> Integer
toInteger a
h
                                  Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= (Int32 -> Integer
forall a. Integral a => a -> Integer
toInteger :: Int32 -> Integer) Int32
forall a. Bounded a => a
maxBound) ((a, SMGen) -> (a, SMGen)) -> (a, SMGen) -> (a, SMGen)
forall a b. (a -> b) -> a -> b
$
  let (Word32
w32, SMGen
g') = Word32 -> SMGen -> (Word32, SMGen)
SM.bitmaskWithRejection32'
                    ((a -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegralWrap :: a -> Word32) a
h) SMGen
g
      -- `fromIntegralWrap` is fine here, because wrapping is OK.
      x :: a
x = (Word32 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegralWrap :: Word32 -> a) Word32
w32
  in if a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
h
     then [Char] -> (a, SMGen)
forall a. HasCallStack => [Char] -> a
error ([Char] -> (a, SMGen)) -> [Char] -> (a, SMGen)
forall a b. (a -> b) -> a -> b
$ [Char]
"nextRandom internal error"
                  [Char] -> (Integer, Integer, Word32) -> [Char]
forall v. Show v => [Char] -> v -> [Char]
`showFailure` (a -> Integer
forall a. Integral a => a -> Integer
toInteger a
x, a -> Integer
forall a. Integral a => a -> Integer
toInteger a
h, Word32
w32)
     else (a
x, SMGen
g')

-- | Get a random 'Word32' using full range.
randomWord32 :: Rnd Word32
{-# INLINE randomWord32 #-}
randomWord32 :: Rnd Word32
randomWord32 = (SMGen -> (Word32, SMGen)) -> Rnd Word32
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
St.state SMGen -> (Word32, SMGen)
SM.nextWord32

-- | 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
$ [Char]
"oneOf []" [Char] -> () -> [Char]
forall v. Show v => [Char] -> v -> [Char]
`showFailure` ()
oneOf [a
x] = a -> Rnd a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x
oneOf [a]
xs = do
  Int
r <- Int -> Rnd Int
forall a. Integral a => a -> Rnd a
randomR0 ([a] -> Int
forall a. [a] -> Int
length [a]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
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 [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)

-- | Code that means the information (e.g., flavour or hidden kind index)
-- should be regenerated, because it could not be transferred from
-- previous playthrough (it's random in each playthrough or there was
-- no previous playthrough).
invalidInformationCode :: Word16
invalidInformationCode :: Word16
invalidInformationCode = Word16
forall a. Bounded a => a
maxBound

-- | Generates a random permutation, except for the existing mapping.
shuffleExcept :: U.Vector Word16 -> Int -> [Word16] -> Rnd [Word16]
shuffleExcept :: Vector Word16 -> Int -> [Word16] -> Rnd [Word16]
shuffleExcept Vector Word16
v Int
len [Word16]
l0 = Bool -> Rnd [Word16] -> Rnd [Word16]
forall a. HasCallStack => Bool -> a -> a
assert (Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Word16] -> Int
forall a. [a] -> Int
length [Word16]
l0) (Rnd [Word16] -> Rnd [Word16]) -> Rnd [Word16] -> Rnd [Word16]
forall a b. (a -> b) -> a -> b
$
  Int -> [Word16] -> Rnd [Word16]
shuffleE Int
0 ([Word16]
l0 [Word16] -> [Word16] -> [Word16]
forall a. Eq a => [a] -> [a] -> [a]
\\ (Word16 -> Bool) -> [Word16] -> [Word16]
forall a. (a -> Bool) -> [a] -> [a]
filter (Word16 -> Word16 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word16
invalidInformationCode) (Vector Word16 -> [Word16]
forall a. Unbox a => Vector a -> [a]
U.toList Vector Word16
v))
 where
  shuffleE :: Int -> [Word16] -> Rnd [Word16]
  shuffleE :: Int -> [Word16] -> Rnd [Word16]
shuffleE Int
i [Word16]
_ | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len = [Word16] -> Rnd [Word16]
forall (m :: * -> *) a. Monad m => a -> m a
return []
  shuffleE Int
i [Word16]
l = do
    let a0 :: Word16
a0 = Vector Word16
v Vector Word16 -> Int -> Word16
forall a. Unbox a => Vector a -> Int -> a
U.! Int
i
    if Word16
a0 Word16 -> Word16 -> Bool
forall a. Eq a => a -> a -> Bool
== Word16
invalidInformationCode then do
      Word16
a <- [Word16] -> Rnd Word16
forall a. [a] -> Rnd a
oneOf [Word16]
l
      (Word16
a Word16 -> [Word16] -> [Word16]
forall a. a -> [a] -> [a]
:) ([Word16] -> [Word16]) -> Rnd [Word16] -> Rnd [Word16]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [Word16] -> Rnd [Word16]
shuffleE (Int -> Int
forall a. Enum a => a -> a
succ Int
i) (Word16 -> [Word16] -> [Word16]
forall a. Eq a => a -> [a] -> [a]
delete Word16
a [Word16]
l)
    else
      (Word16
a0 Word16 -> [Word16] -> [Word16]
forall a. a -> [a] -> [a]
:) ([Word16] -> [Word16]) -> Rnd [Word16] -> Rnd [Word16]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [Word16] -> Rnd [Word16]
shuffleE (Int -> Int
forall a. Enum a => a -> a
succ Int
i) [Word16]
l

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

-- | Randomly choose an item according to the distribution.
rollFreq :: Show a => Frequency a -> SM.SMGen -> (a, SM.SMGen)
rollFreq :: Frequency a -> SMGen -> (a, SMGen)
rollFreq Frequency a
fr SMGen
g = case Frequency a -> [(Int, a)]
forall a. Frequency a -> [(Int, a)]
runFrequency Frequency a
fr of
  [] -> [Char] -> (a, SMGen)
forall a. HasCallStack => [Char] -> a
error ([Char] -> (a, SMGen)) -> [Char] -> (a, SMGen)
forall a b. (a -> b) -> a -> b
$ [Char]
"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
  [(Int
n, a
x)] | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 -> [Char] -> (a, SMGen)
forall a. HasCallStack => [Char] -> a
error ([Char] -> (a, SMGen)) -> [Char] -> (a, SMGen)
forall a b. (a -> b) -> a -> b
$ [Char]
"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)
  [(Int
_, a
x)] -> (a
x, SMGen
g)  -- speedup
  [(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, a
_) -> Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) Int
0 [(Int, a)]
fs
            (Int
r, SMGen
ng) = Int -> SMGen -> (Int, SMGen)
forall a. Integral a => a -> SMGen -> (a, SMGen)
nextRandom (Int -> Int
forall a. Enum a => a -> a
pred Int
sumf) SMGen
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
$ [Char]
"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 Int
m ((Int
n, a
x) : [(Int, a)]
_) | Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n = a
x
            frec Int
m ((Int
n, a
_) : [(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, SMGen) -> (a, SMGen)
forall a. HasCallStack => Bool -> a -> a
assert (Int
sumf Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> ([Char], (Text, [(Int, a)])) -> Bool
forall a. Show a => Bool -> a -> Bool
`blame` [Char]
"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, SMGen
ng)

-- | Fractional chance.
type Chance = Rational

-- | Give @True@, with probability determined by the fraction.
chance :: Chance -> Rnd Bool
chance :: Chance -> Rnd Bool
chance 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. Integral a => (a, a) -> Rnd a
randomR (Integer
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.
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. Integral 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 AbsDepth
ldepth AbsDepth
totalDepth 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
> Int
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 AbsDepth
ldepth AbsDepth
totalDepth (Dice.DiceXY Dice
dx 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 a -> b -> Rnd b
f b
z0 t a
xs = let f' :: a -> (b, SMGen) -> (b, SMGen)
f' a
x (b
z, SMGen
g) = Rnd b -> SMGen -> (b, SMGen)
forall s a. State s a -> s -> (a, s)
St.runState (a -> b -> Rnd b
f a
x b
z) SMGen
g
                 in (SMGen -> (b, SMGen)) -> Rnd b
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
St.state ((SMGen -> (b, SMGen)) -> Rnd b) -> (SMGen -> (b, SMGen)) -> Rnd b
forall a b. (a -> b) -> a -> b
$ \SMGen
g -> (a -> (b, SMGen) -> (b, SMGen)) -> (b, SMGen) -> t a -> (b, SMGen)
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> (b, SMGen) -> (b, SMGen)
f' (b
z0, SMGen
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' b -> a -> Rnd b
f b
z0 t a
xs = let f' :: (b, SMGen) -> a -> (b, SMGen)
f' (b
z, SMGen
g) a
x = Rnd b -> SMGen -> (b, SMGen)
forall s a. State s a -> s -> (a, s)
St.runState (b -> a -> Rnd b
f b
z a
x) SMGen
g
                  in (SMGen -> (b, SMGen)) -> Rnd b
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
St.state ((SMGen -> (b, SMGen)) -> Rnd b) -> (SMGen -> (b, SMGen)) -> Rnd b
forall a b. (a -> b) -> a -> b
$ \SMGen
g -> ((b, SMGen) -> a -> (b, SMGen)) -> (b, SMGen) -> t a -> (b, SMGen)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (b, SMGen) -> a -> (b, SMGen)
f' (b
z0, SMGen
g) t a
xs