module Game.LambdaHack.Utils.Frequency
(
Frequency
, uniformFreq, toFreq
, scaleFreq, renameFreq, setFreq
, rollFreq, nullFreq, runFrequency, nameFrequency
) where
import Control.Applicative
import Control.Arrow (first, second)
import Control.Monad
import Data.Binary
import Data.Foldable (Foldable)
import Data.Text (Text)
import Data.Traversable (Traversable)
import qualified System.Random as R
import Control.Exception.Assert.Sugar
import Game.LambdaHack.Common.Msg
data Frequency a = Frequency
{ nameFrequency :: !Text
, runFrequency :: ![(Int, a)]
}
deriving (Show, Eq, Foldable, Traversable)
instance Monad Frequency where
return x = Frequency "return" [(1, x)]
Frequency name xs >>= f =
Frequency ("bind (" <> name <> ")")
[(p * q, y) | (p, x) <- xs
, (q, y) <- runFrequency (f x) ]
instance Functor Frequency where
fmap f (Frequency name xs) = Frequency name (map (second f) xs)
instance Applicative Frequency where
pure = return
(<*>) = ap
instance MonadPlus Frequency where
mplus (Frequency xname xs) (Frequency yname ys) =
let name = case (xs, ys) of
([], []) -> "[]"
([], _ ) -> yname
(_, []) -> xname
_ -> "(" <> xname <> ") ++ (" <> yname <> ")"
in Frequency name (xs ++ ys)
mzero = Frequency "[]" []
instance Alternative Frequency where
(<|>) = mplus
empty = mzero
uniformFreq :: Text -> [a] -> Frequency a
uniformFreq name = Frequency name . map (\ x -> (1, x))
toFreq :: Text -> [(Int, a)] -> Frequency a
toFreq = Frequency
scaleFreq :: Show a => Int -> Frequency a -> Frequency a
scaleFreq n (Frequency name xs) =
assert (n > 0 `blame` "non-positive frequency scale" `twith` (name, n, xs)) $
Frequency name (map (first (* n)) xs)
renameFreq :: Text -> Frequency a -> Frequency a
renameFreq newName fr = fr {nameFrequency = newName}
setFreq :: Eq a => Frequency a -> a -> Int -> Frequency a
setFreq (Frequency name xs) x n =
let f (_, y) | y == x = (n, x)
f my = my
in Frequency name $ map f xs
rollFreq :: Show a => Frequency a -> R.StdGen -> (a, R.StdGen)
rollFreq (Frequency name []) _ =
assert `failure` "choice from an empty frequency" `twith` name
rollFreq (Frequency name [(n, x)]) _ | n <= 0 =
assert `failure` "singleton void frequency" `twith` (name, n, x)
rollFreq (Frequency _ [(_, x)]) g = (x, g)
rollFreq (Frequency name fs) g =
assert (sumf > 0 `blame` "frequency with nothing to pick" `twith` (name, fs))
(frec r fs, ng)
where
sumf = sum (map fst fs)
(r, ng) = R.randomR (1, sumf) g
frec :: Int -> [(Int, a)] -> a
frec m [] = assert `failure` "impossible roll"
`twith` (name, fs, m)
frec m ((n, x) : _) | m <= n = x
frec m ((n, _) : xs) = frec (m n) xs
nullFreq :: Frequency a -> Bool
nullFreq (Frequency _ fs) = all (== 0) $ map fst fs
instance Binary a => Binary (Frequency a) where
put Frequency{..} = do
put nameFrequency
put runFrequency
get = do
nameFrequency <- get
runFrequency <- get
return Frequency{..}