-- |
-- Module      : MonusWeightedSearch.Examples.Sampling
-- Copyright   : (c) Donnacha Oisín Kidney 2021
-- Maintainer  : mail@doisinkidney.com
-- Stability   : experimental
-- Portability : non-portable
--
-- Random sampling from the 'Heap' monad.
--
-- The 'Heap' monad can function as a probability monad, and it implements an
-- efficient sampling algorithm, based on reservoir sampling.

module MonusWeightedSearch.Examples.Sampling where

import Control.Monad.Heap
import Data.Monus.Prob
import Data.Ratio
import System.Random (randomRIO)

-- | @'withChance' p@ returns 'True' @p@ percent of the time.
withChance :: Integral a => Ratio a -> IO Bool
withChance :: forall a. Integral a => Ratio a -> IO Bool
withChance Ratio a
f = (Integer -> Bool) -> IO Integer -> IO Bool
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (a -> Integer
forall a. Integral a => a -> Integer
toInteger (Ratio a -> a
forall a. Ratio a -> a
numerator Ratio a
f) Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>=) ((Integer, Integer) -> IO Integer
forall a (m :: Type -> Type).
(Random a, MonadIO m) =>
(a, a) -> m a
randomRIO (Integer
1, a -> Integer
forall a. Integral a => a -> Integer
toInteger (Ratio a -> a
forall a. Ratio a -> a
denominator Ratio a
f)))

-- | Sample a single value from the heap.
sample :: Heap Prob a -> IO a
sample :: forall a. Heap Prob a -> IO a
sample Heap Prob a
hp = ((a, Prob) -> (Ratio Natural -> IO a) -> Ratio Natural -> IO a)
-> (Ratio Natural -> IO a) -> [(a, Prob)] -> Ratio Natural -> IO a
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (a, Prob) -> (Ratio Natural -> IO a) -> Ratio Natural -> IO a
forall {b}.
(b, Prob) -> (Ratio Natural -> IO b) -> Ratio Natural -> IO b
f ([Char] -> Ratio Natural -> IO a
forall a. HasCallStack => [Char] -> a
error [Char]
"impossible") (Heap Prob a -> [(a, Prob)]
forall w a. Monus w => Heap w a -> [(a, w)]
search Heap Prob a
hp) Ratio Natural
1 where
  f :: (b, Prob) -> (Ratio Natural -> IO b) -> Ratio Natural -> IO b
f (b
x,Prob Ratio Natural
px) Ratio Natural -> IO b
k Ratio Natural
r = do
    let f :: Ratio Natural
f = Ratio Natural
r Ratio Natural -> Ratio Natural -> Ratio Natural
forall a. Num a => a -> a -> a
* Ratio Natural
px
    Bool
c <- Ratio Natural -> IO Bool
forall a. Integral a => Ratio a -> IO Bool
withChance Ratio Natural
f
    if Bool
c then b -> IO b
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure b
x else Ratio Natural -> IO b
k (Ratio Natural
r Ratio Natural -> Ratio Natural -> Ratio Natural
forall a. Fractional a => a -> a -> a
/ (Ratio Natural
1 Ratio Natural -> Ratio Natural -> Ratio Natural
forall a. Num a => a -> a -> a
- Ratio Natural
f))