-- Module      : Verismith.Verilog2005.Randomness
-- Description : Random sources
-- Copyright   : (c) 2023 Quentin Corradi
-- License     : GPL-3
-- Maintainer  : q [dot] corradi22 [at] imperial [dot] ac [dot] uk
-- Stability   : stable
-- Portability : POSIX
{-# LANGUAGE ScopedTypeVariables #-}

module Verismith.Verilog2005.Randomness
  ( sampleCategoricalProbability,
    sampleNumberProbability,
    sampleIn,
    sampleInString,
    sampleBernoulli,
    choice,
    sampleMaybe,
    sampleEither,
    sampleSegment,
    sampleEnum,
    sampleMaybeEnum,
    sampleWeighted,
    sampleFrom,
    sampleFromString,
    sampleBranch,
    sampleNum,
    sampleN,
    sampleNE,
    sampleString,
    sampleNEString,
    sampleFiltered,
    GenM,
  )
where

import Control.Applicative (liftA2)
import Control.Monad (join, replicateM)
import Control.Monad.Reader
import qualified Data.ByteString as B
import Data.List
import Data.List.NonEmpty (NonEmpty (..), toList)
import qualified Data.List.NonEmpty as NE
import Control.Monad.Primitive (PrimMonad, PrimState, RealWorld)
import Data.Word
import System.Random.MWC.Probability
import Verismith.Config (CategoricalProbability (..), NumberProbability (..), uniformCP)
import Verismith.Utils (nonEmpty, foldrMap1)

infixl 4 <.>

(<.>) :: (Monad m, Applicative m) => m (a -> m b) -> m a -> m b
<.> :: forall (m :: * -> *) a b.
(Monad m, Applicative m) =>
m (a -> m b) -> m a -> m b
(<.>) m (a -> m b)
mf m a
mx = m (m b) -> m b
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (m (m b) -> m b) -> m (m b) -> m b
forall a b. (a -> b) -> a -> b
$ m (a -> m b)
mf m (a -> m b) -> m a -> m (m b)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> m a
mx

avoid :: [Int] -> Int -> Int
avoid :: [Int] -> Int -> Int
avoid [Int]
l Int
x = case [Int]
l of
  Int
h : [Int]
t | Int
h Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
x -> [Int] -> Int -> Int
avoid [Int]
t (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
  [Int]
_ -> Int
x

uniq :: Ord b => (a -> b) -> (a -> a -> a) -> [a] -> [a]
uniq :: forall b a. Ord b => (a -> b) -> (a -> a -> a) -> [a] -> [a]
uniq a -> b
f a -> a -> a
m =
  [a] -> (NonEmpty a -> [a]) -> [a] -> [a]
forall b a. b -> (NonEmpty a -> b) -> [a] -> b
nonEmpty [] ((NonEmpty a -> [a]) -> [a] -> [a])
-> (NonEmpty a -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
toList
    (NonEmpty a -> [a])
-> (NonEmpty a -> NonEmpty a) -> NonEmpty a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> NonEmpty a)
-> (a -> NonEmpty a -> NonEmpty a) -> NonEmpty a -> NonEmpty a
forall a b. (a -> b) -> (a -> b -> b) -> NonEmpty a -> b
foldrMap1 (a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:|[]) (\a
e (a
x :| [a]
a) -> if a -> b
f a
x b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== a -> b
f a
e then (a -> a -> a
m a
x a
e) a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
a else a
e a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
a)
    (NonEmpty a -> NonEmpty a)
-> (NonEmpty a -> NonEmpty a) -> NonEmpty a -> NonEmpty a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b) -> NonEmpty a -> NonEmpty a
forall o a. Ord o => (a -> o) -> NonEmpty a -> NonEmpty a
NE.sortWith a -> b
f
    

clean :: Int -> [(Double, Int)] -> [(Double, Int)]
clean :: Int -> [(Double, Int)] -> [(Double, Int)]
clean Int
t =
  ((Double, Int) -> (Double, Int))
-> [(Double, Int)] -> [(Double, Int)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Double
x, Int
y) -> (Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0 Double
x, Int
y))
    ([(Double, Int)] -> [(Double, Int)])
-> ([(Double, Int)] -> [(Double, Int)])
-> [(Double, Int)]
-> [(Double, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Double, Int) -> Int)
-> ((Double, Int) -> (Double, Int) -> (Double, Int))
-> [(Double, Int)]
-> [(Double, Int)]
forall b a. Ord b => (a -> b) -> (a -> a -> a) -> [a] -> [a]
uniq (Double, Int) -> Int
forall a b. (a, b) -> b
snd (\(Double
x1, Int
y1) (Double
x2, Int
y2) -> (Double
x1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
x2, Int
y1))
    ([(Double, Int)] -> [(Double, Int)])
-> ([(Double, Int)] -> [(Double, Int)])
-> [(Double, Int)]
-> [(Double, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Double, Int) -> Bool) -> [(Double, Int)] -> [(Double, Int)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
t) (Int -> Bool) -> ((Double, Int) -> Int) -> (Double, Int) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double, Int) -> Int
forall a b. (a, b) -> b
snd)

sampleCategoricalProbability ::
  PrimMonad m => Int -> Gen (PrimState m) -> CategoricalProbability -> m Int
sampleCategoricalProbability :: forall (m :: * -> *).
PrimMonad m =>
Int -> Gen (PrimState m) -> CategoricalProbability -> m Int
sampleCategoricalProbability Int
t Gen (PrimState m)
gen CategoricalProbability
d = case CategoricalProbability
d of
  CPDiscrete NonEmpty Double
l ->
    let ll :: [Double]
ll = Int -> NonEmpty Double -> [Double]
forall a. Int -> NonEmpty a -> [a]
NE.take (Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) NonEmpty Double
l
     in case [Double]
ll of
          [] -> [Char] -> m Int
forall a. HasCallStack => [Char] -> a
error [Char]
"Probability vector cannot be empty"
          [Double
x] -> Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
0
          [Double]
_ -> Prob m Int -> Gen (PrimState m) -> m Int
forall (m :: * -> *) a. Prob m a -> Gen (PrimState m) -> m a
sample ([Double] -> Prob m Int
forall (f :: * -> *) (m :: * -> *).
(Foldable f, PrimMonad m) =>
f Double -> Prob m Int
categorical [Double]
ll) Gen (PrimState m)
gen
  CPBiasedUniform [(Double, Int)]
l Double
b ->
    let ll :: [(Double, Int)]
ll = Int -> [(Double, Int)] -> [(Double, Int)]
clean Int
t [(Double, Int)]
l
        uw :: Double
uw = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- [(Double, Int)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Double, Int)]
ll) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b
     in m (Maybe Int)
-> (NonEmpty (Double, Int) -> m (Maybe Int))
-> [(Double, Int)]
-> m (Maybe Int)
forall b a. b -> (NonEmpty a -> b) -> [a] -> b
nonEmpty
          (Maybe Int -> m (Maybe Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
Nothing)
          ((Prob m (Maybe Int) -> Gen (PrimState m) -> m (Maybe Int))
-> Gen (PrimState m) -> Prob m (Maybe Int) -> m (Maybe Int)
forall a b c. (a -> b -> c) -> b -> a -> c
flip Prob m (Maybe Int) -> Gen (PrimState m) -> m (Maybe Int)
forall (m :: * -> *) a. Prob m a -> Gen (PrimState m) -> m a
sample Gen (PrimState m)
gen (Prob m (Maybe Int) -> m (Maybe Int))
-> (NonEmpty (Double, Int) -> Prob m (Maybe Int))
-> NonEmpty (Double, Int)
-> m (Maybe Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Double, Maybe Int)] -> Prob m (Maybe Int)
forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, PrimMonad m) =>
f (Double, a) -> Prob m a
discrete ([(Double, Maybe Int)] -> Prob m (Maybe Int))
-> (NonEmpty (Double, Int) -> [(Double, Maybe Int)])
-> NonEmpty (Double, Int)
-> Prob m (Maybe Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Double
uw, Maybe Int
forall a. Maybe a
Nothing) (Double, Maybe Int)
-> [(Double, Maybe Int)] -> [(Double, Maybe Int)]
forall a. a -> [a] -> [a]
:) ([(Double, Maybe Int)] -> [(Double, Maybe Int)])
-> (NonEmpty (Double, Int) -> [(Double, Maybe Int)])
-> NonEmpty (Double, Int)
-> [(Double, Maybe Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Double, Int) -> (Double, Maybe Int))
-> [(Double, Int)] -> [(Double, Maybe Int)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Double
x, Int
y) -> (Double
x, Int -> Maybe Int
forall a. a -> Maybe a
Just Int
y)) ([(Double, Int)] -> [(Double, Maybe Int)])
-> (NonEmpty (Double, Int) -> [(Double, Int)])
-> NonEmpty (Double, Int)
-> [(Double, Maybe Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty (Double, Int) -> [(Double, Int)]
forall a. NonEmpty a -> [a]
toList)
          [(Double, Int)]
ll
          m (Maybe Int) -> (Maybe Int -> m Int) -> m Int
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= m Int -> (Int -> m Int) -> Maybe Int -> m Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ([Int] -> Int -> Int
avoid (((Double, Int) -> Int) -> [(Double, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Double, Int) -> Int
forall a b. (a, b) -> b
snd [(Double, Int)]
ll) (Int -> Int) -> m Int -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Prob m Int -> Gen (PrimState m) -> m Int
forall (m :: * -> *) a. Prob m a -> Gen (PrimState m) -> m a
sample ((Int, Int) -> Prob m Int
forall (m :: * -> *) a.
(PrimMonad m, Variate a) =>
(a, a) -> Prob m a
uniformR (Int
0, Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
- [(Double, Int)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Double, Int)]
ll)) Gen (PrimState m)
gen) Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

sampleNumberProbability :: PrimMonad m => Gen (PrimState m) -> NumberProbability -> m Int
sampleNumberProbability :: forall (m :: * -> *).
PrimMonad m =>
Gen (PrimState m) -> NumberProbability -> m Int
sampleNumberProbability Gen (PrimState m)
gen NumberProbability
d = case NumberProbability
d of
  NPUniform Int
l Int
h -> Prob m Int -> Gen (PrimState m) -> m Int
forall (m :: * -> *) a. Prob m a -> Gen (PrimState m) -> m a
sample ((Int, Int) -> Prob m Int
forall (m :: * -> *) a.
(PrimMonad m, Variate a) =>
(a, a) -> Prob m a
uniformR (Int
l, Int
h)) Gen (PrimState m)
gen
  NPBinomial Int
o Int
t Double
f -> (Int
o Int -> Int -> Int
forall a. Num a => a -> a -> a
+) (Int -> Int) -> m Int -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Prob m Int -> Gen (PrimState m) -> m Int
forall (m :: * -> *) a. Prob m a -> Gen (PrimState m) -> m a
sample (Int -> Double -> Prob m Int
forall (m :: * -> *). PrimMonad m => Int -> Double -> Prob m Int
binomial Int
t Double
f) Gen (PrimState m)
gen
  NPNegativeBinomial Int
o Double
r Int
f -> (Int
o Int -> Int -> Int
forall a. Num a => a -> a -> a
+) (Int -> Int) -> m Int -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Prob m Int -> Gen (PrimState m) -> m Int
forall (m :: * -> *) a. Prob m a -> Gen (PrimState m) -> m a
sample (Int -> Double -> Prob m Int
forall (m :: * -> *) a.
(PrimMonad m, Integral a) =>
a -> Double -> Prob m Int
negativeBinomial Int
f Double
r) Gen (PrimState m)
gen
  NPPoisson Int
o Double
p -> (Int
o Int -> Int -> Int
forall a. Num a => a -> a -> a
+) (Int -> Int) -> m Int -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Prob m Int -> Gen (PrimState m) -> m Int
forall (m :: * -> *) a. Prob m a -> Gen (PrimState m) -> m a
sample (Double -> Prob m Int
forall (m :: * -> *). PrimMonad m => Double -> Prob m Int
poisson Double
p) Gen (PrimState m)
gen
  NPDiscrete NonEmpty (Double, Int)
l -> Prob m Int -> Gen (PrimState m) -> m Int
forall (m :: * -> *) a. Prob m a -> Gen (PrimState m) -> m a
sample (NonEmpty (Double, Int) -> Prob m Int
forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, PrimMonad m) =>
f (Double, a) -> Prob m a
discrete NonEmpty (Double, Int)
l) Gen (PrimState m)
gen
  NPLinearComb NonEmpty (Double, NumberProbability)
l -> Prob m NumberProbability
-> Gen (PrimState m) -> m NumberProbability
forall (m :: * -> *) a. Prob m a -> Gen (PrimState m) -> m a
sample (NonEmpty (Double, NumberProbability) -> Prob m NumberProbability
forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, PrimMonad m) =>
f (Double, a) -> Prob m a
discrete NonEmpty (Double, NumberProbability)
l) Gen (PrimState m)
gen m NumberProbability -> (NumberProbability -> m Int) -> m Int
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Gen (PrimState m) -> NumberProbability -> m Int
forall (m :: * -> *).
PrimMonad m =>
Gen (PrimState m) -> NumberProbability -> m Int
sampleNumberProbability Gen (PrimState m)
gen

sampleIn :: (Functor m, PrimMonad m) => [a] -> Gen (PrimState m) -> CategoricalProbability -> m a
sampleIn :: forall (m :: * -> *) a.
(Functor m, PrimMonad m) =>
[a] -> Gen (PrimState m) -> CategoricalProbability -> m a
sampleIn [a]
l Gen (PrimState m)
gen CategoricalProbability
d = ([a]
l [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!!) (Int -> a) -> m Int -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Gen (PrimState m) -> CategoricalProbability -> m Int
forall (m :: * -> *).
PrimMonad m =>
Int -> Gen (PrimState m) -> CategoricalProbability -> m Int
sampleCategoricalProbability ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Gen (PrimState m)
gen CategoricalProbability
d

sampleInString ::
  (Functor m, PrimMonad m) =>
  B.ByteString ->
  Gen (PrimState m) ->
  CategoricalProbability ->
  m Word8
sampleInString :: forall (m :: * -> *).
(Functor m, PrimMonad m) =>
ByteString
-> Gen (PrimState m) -> CategoricalProbability -> m Word8
sampleInString ByteString
s Gen (PrimState m)
gen CategoricalProbability
d = (HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
B.index ByteString
s) (Int -> Word8) -> m Int -> m Word8
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Gen (PrimState m) -> CategoricalProbability -> m Int
forall (m :: * -> *).
PrimMonad m =>
Int -> Gen (PrimState m) -> CategoricalProbability -> m Int
sampleCategoricalProbability (ByteString -> Int
B.length ByteString
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Gen (PrimState m)
gen CategoricalProbability
d

type GenM p = ReaderT (p, Gen RealWorld) IO

sampleWrapper :: (p -> d) -> (Gen RealWorld -> d -> GenM p x) -> GenM p x
sampleWrapper :: forall p d x.
(p -> d) -> (Gen RealWorld -> d -> GenM p x) -> GenM p x
sampleWrapper p -> d
p Gen RealWorld -> d -> GenM p x
f = Gen RealWorld -> d -> GenM p x
f (Gen RealWorld -> d -> GenM p x)
-> ReaderT (p, Gen RealWorld) IO (Gen RealWorld)
-> ReaderT (p, Gen RealWorld) IO (d -> GenM p x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((p, Gen RealWorld) -> Gen RealWorld)
-> ReaderT (p, Gen RealWorld) IO (Gen RealWorld)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (p, Gen RealWorld) -> Gen RealWorld
forall a b. (a, b) -> b
snd ReaderT (p, Gen RealWorld) IO (d -> GenM p x)
-> ReaderT (p, Gen RealWorld) IO d -> GenM p x
forall (m :: * -> *) a b.
(Monad m, Applicative m) =>
m (a -> m b) -> m a -> m b
<.> ((p, Gen RealWorld) -> d) -> ReaderT (p, Gen RealWorld) IO d
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (p -> d
p (p -> d) -> ((p, Gen RealWorld) -> p) -> (p, Gen RealWorld) -> d
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (p, Gen RealWorld) -> p
forall a b. (a, b) -> a
fst)

sampleBernoulli :: (p -> Double) -> GenM p Bool
sampleBernoulli :: forall p. (p -> Double) -> GenM p Bool
sampleBernoulli p -> Double
p = Prob (ReaderT (p, Gen RealWorld) IO) Bool
-> Gen RealWorld -> ReaderT (p, Gen RealWorld) IO Bool
Prob (ReaderT (p, Gen RealWorld) IO) Bool
-> Gen (PrimState (ReaderT (p, Gen RealWorld) IO))
-> ReaderT (p, Gen RealWorld) IO Bool
forall (m :: * -> *) a. Prob m a -> Gen (PrimState m) -> m a
sample (Prob (ReaderT (p, Gen RealWorld) IO) Bool
 -> Gen RealWorld -> ReaderT (p, Gen RealWorld) IO Bool)
-> ReaderT
     (p, Gen RealWorld) IO (Prob (ReaderT (p, Gen RealWorld) IO) Bool)
-> ReaderT
     (p, Gen RealWorld)
     IO
     (Gen RealWorld -> ReaderT (p, Gen RealWorld) IO Bool)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Double -> Prob (ReaderT (p, Gen RealWorld) IO) Bool
forall (m :: * -> *). PrimMonad m => Double -> Prob m Bool
bernoulli (Double -> Prob (ReaderT (p, Gen RealWorld) IO) Bool)
-> ReaderT (p, Gen RealWorld) IO Double
-> ReaderT
     (p, Gen RealWorld) IO (Prob (ReaderT (p, Gen RealWorld) IO) Bool)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((p, Gen RealWorld) -> Double)
-> ReaderT (p, Gen RealWorld) IO Double
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (p -> Double
p (p -> Double)
-> ((p, Gen RealWorld) -> p) -> (p, Gen RealWorld) -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (p, Gen RealWorld) -> p
forall a b. (a, b) -> a
fst)) ReaderT
  (p, Gen RealWorld)
  IO
  (Gen RealWorld -> ReaderT (p, Gen RealWorld) IO Bool)
-> ReaderT (p, Gen RealWorld) IO (Gen RealWorld)
-> ReaderT (p, Gen RealWorld) IO Bool
forall (m :: * -> *) a b.
(Monad m, Applicative m) =>
m (a -> m b) -> m a -> m b
<.> ((p, Gen RealWorld) -> Gen RealWorld)
-> ReaderT (p, Gen RealWorld) IO (Gen RealWorld)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (p, Gen RealWorld) -> Gen RealWorld
forall a b. (a, b) -> b
snd

choice :: (p -> Double) -> GenM p a -> GenM p a -> GenM p a
choice :: forall p a. (p -> Double) -> GenM p a -> GenM p a -> GenM p a
choice p -> Double
c GenM p a
t GenM p a
f = (p -> Double) -> GenM p Bool
forall p. (p -> Double) -> GenM p Bool
sampleBernoulli p -> Double
c GenM p Bool -> (Bool -> GenM p a) -> GenM p a
forall a b.
ReaderT (p, Gen RealWorld) IO a
-> (a -> ReaderT (p, Gen RealWorld) IO b)
-> ReaderT (p, Gen RealWorld) IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Bool
b -> if Bool
b then GenM p a
t else GenM p a
f

sampleMaybe :: (p -> Double) -> GenM p a -> GenM p (Maybe a)
sampleMaybe :: forall p a. (p -> Double) -> GenM p a -> GenM p (Maybe a)
sampleMaybe p -> Double
c GenM p a
x = (p -> Double)
-> GenM p (Maybe a) -> GenM p (Maybe a) -> GenM p (Maybe a)
forall p a. (p -> Double) -> GenM p a -> GenM p a -> GenM p a
choice p -> Double
c (a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> GenM p a -> GenM p (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GenM p a
x) (Maybe a -> GenM p (Maybe a)
forall a. a -> ReaderT (p, Gen RealWorld) IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing)

sampleEither :: (p -> Double) -> GenM p a -> GenM p b -> GenM p (Either a b)
sampleEither :: forall p a b.
(p -> Double) -> GenM p a -> GenM p b -> GenM p (Either a b)
sampleEither p -> Double
c GenM p a
t GenM p b
f = (p -> Double)
-> GenM p (Either a b)
-> GenM p (Either a b)
-> GenM p (Either a b)
forall p a. (p -> Double) -> GenM p a -> GenM p a -> GenM p a
choice p -> Double
c (a -> Either a b
forall a b. a -> Either a b
Left (a -> Either a b) -> GenM p a -> GenM p (Either a b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GenM p a
t) (b -> Either a b
forall a b. b -> Either a b
Right (b -> Either a b) -> GenM p b -> GenM p (Either a b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GenM p b
f)

sampleSegment :: (p -> CategoricalProbability) -> Int -> Int -> GenM p Int
sampleSegment :: forall p. (p -> CategoricalProbability) -> Int -> Int -> GenM p Int
sampleSegment p -> CategoricalProbability
p Int
l Int
h = (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+) (Int -> Int)
-> ReaderT (p, Gen RealWorld) IO Int
-> ReaderT (p, Gen RealWorld) IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((p -> CategoricalProbability)
-> (Gen RealWorld
    -> CategoricalProbability -> ReaderT (p, Gen RealWorld) IO Int)
-> ReaderT (p, Gen RealWorld) IO Int
forall p d x.
(p -> d) -> (Gen RealWorld -> d -> GenM p x) -> GenM p x
sampleWrapper p -> CategoricalProbability
p ((Gen RealWorld
  -> CategoricalProbability -> ReaderT (p, Gen RealWorld) IO Int)
 -> ReaderT (p, Gen RealWorld) IO Int)
-> (Gen RealWorld
    -> CategoricalProbability -> ReaderT (p, Gen RealWorld) IO Int)
-> ReaderT (p, Gen RealWorld) IO Int
forall a b. (a -> b) -> a -> b
$ Int
-> Gen (PrimState (ReaderT (p, Gen RealWorld) IO))
-> CategoricalProbability
-> ReaderT (p, Gen RealWorld) IO Int
forall (m :: * -> *).
PrimMonad m =>
Int -> Gen (PrimState m) -> CategoricalProbability -> m Int
sampleCategoricalProbability (Int
 -> Gen (PrimState (ReaderT (p, Gen RealWorld) IO))
 -> CategoricalProbability
 -> ReaderT (p, Gen RealWorld) IO Int)
-> Int
-> Gen (PrimState (ReaderT (p, Gen RealWorld) IO))
-> CategoricalProbability
-> ReaderT (p, Gen RealWorld) IO Int
forall a b. (a -> b) -> a -> b
$ Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l)

sampleEnum :: forall a p. (Bounded a, Enum a) => (p -> CategoricalProbability) -> GenM p a
sampleEnum :: forall a p.
(Bounded a, Enum a) =>
(p -> CategoricalProbability) -> GenM p a
sampleEnum p -> CategoricalProbability
p = Int -> a
forall a. Enum a => Int -> a
toEnum (Int -> a)
-> ReaderT (p, Gen RealWorld) IO Int
-> ReaderT (p, Gen RealWorld) IO a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (p -> CategoricalProbability)
-> Int -> Int -> ReaderT (p, Gen RealWorld) IO Int
forall p. (p -> CategoricalProbability) -> Int -> Int -> GenM p Int
sampleSegment p -> CategoricalProbability
p (a -> Int
forall a. Enum a => a -> Int
fromEnum (a
forall a. Bounded a => a
minBound :: a)) (a -> Int
forall a. Enum a => a -> Int
fromEnum (a
forall a. Bounded a => a
maxBound :: a))

sampleMaybeEnum ::
  forall a p. (Bounded a, Enum a) => (p -> CategoricalProbability) -> GenM p (Maybe a)
sampleMaybeEnum :: forall a p.
(Bounded a, Enum a) =>
(p -> CategoricalProbability) -> GenM p (Maybe a)
sampleMaybeEnum p -> CategoricalProbability
p =
  (\Int
n -> if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Maybe a
forall a. Maybe a
Nothing else a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ Int -> a
forall a. Enum a => Int -> a
toEnum (Int -> a) -> Int -> a
forall a b. (a -> b) -> a -> b
$ Int
mib Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
    (Int -> Maybe a)
-> ReaderT (p, Gen RealWorld) IO Int
-> ReaderT (p, Gen RealWorld) IO (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((p -> CategoricalProbability)
-> (Gen RealWorld
    -> CategoricalProbability -> ReaderT (p, Gen RealWorld) IO Int)
-> ReaderT (p, Gen RealWorld) IO Int
forall p d x.
(p -> d) -> (Gen RealWorld -> d -> GenM p x) -> GenM p x
sampleWrapper p -> CategoricalProbability
p ((Gen RealWorld
  -> CategoricalProbability -> ReaderT (p, Gen RealWorld) IO Int)
 -> ReaderT (p, Gen RealWorld) IO Int)
-> (Gen RealWorld
    -> CategoricalProbability -> ReaderT (p, Gen RealWorld) IO Int)
-> ReaderT (p, Gen RealWorld) IO Int
forall a b. (a -> b) -> a -> b
$ Int
-> Gen (PrimState (ReaderT (p, Gen RealWorld) IO))
-> CategoricalProbability
-> ReaderT (p, Gen RealWorld) IO Int
forall (m :: * -> *).
PrimMonad m =>
Int -> Gen (PrimState m) -> CategoricalProbability -> m Int
sampleCategoricalProbability (Int
 -> Gen (PrimState (ReaderT (p, Gen RealWorld) IO))
 -> CategoricalProbability
 -> ReaderT (p, Gen RealWorld) IO Int)
-> Int
-> Gen (PrimState (ReaderT (p, Gen RealWorld) IO))
-> CategoricalProbability
-> ReaderT (p, Gen RealWorld) IO Int
forall a b. (a -> b) -> a -> b
$ Int
mab Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
mib Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  where
    mib :: Int
mib = a -> Int
forall a. Enum a => a -> Int
fromEnum (a
forall a. Bounded a => a
minBound :: a)
    mab :: Int
mab = a -> Int
forall a. Enum a => a -> Int
fromEnum (a
forall a. Bounded a => a
maxBound :: a)

sampleWeighted :: [(Double, a)] -> GenM p a
sampleWeighted :: forall a p. [(Double, a)] -> GenM p a
sampleWeighted [(Double, a)]
l = case [(Double, a)]
l of
  [] -> [Char] -> GenM p a
forall a. HasCallStack => [Char] -> a
error [Char]
"Probability vector cannot be empty"
  [(Double
_, a
x)] -> a -> GenM p a
forall a. a -> ReaderT (p, Gen RealWorld) IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x
  [(Double, a)]
_ -> ((p, Gen RealWorld) -> Gen RealWorld)
-> ReaderT (p, Gen RealWorld) IO (Gen RealWorld)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (p, Gen RealWorld) -> Gen RealWorld
forall a b. (a, b) -> b
snd ReaderT (p, Gen RealWorld) IO (Gen RealWorld)
-> (Gen RealWorld -> GenM p a) -> GenM p a
forall a b.
ReaderT (p, Gen RealWorld) IO a
-> (a -> ReaderT (p, Gen RealWorld) IO b)
-> ReaderT (p, Gen RealWorld) IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Prob (ReaderT (p, Gen RealWorld) IO) a
-> Gen (PrimState (ReaderT (p, Gen RealWorld) IO)) -> GenM p a
forall (m :: * -> *) a. Prob m a -> Gen (PrimState m) -> m a
sample ([(Double, a)] -> Prob (ReaderT (p, Gen RealWorld) IO) a
forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, PrimMonad m) =>
f (Double, a) -> Prob m a
discrete [(Double, a)]
l)

sampleFrom :: (p -> CategoricalProbability) -> [a] -> GenM p a
sampleFrom :: forall p a. (p -> CategoricalProbability) -> [a] -> GenM p a
sampleFrom p -> CategoricalProbability
p [a]
l = (p -> CategoricalProbability)
-> (Gen RealWorld -> CategoricalProbability -> GenM p a)
-> GenM p a
forall p d x.
(p -> d) -> (Gen RealWorld -> d -> GenM p x) -> GenM p x
sampleWrapper p -> CategoricalProbability
p ((Gen RealWorld -> CategoricalProbability -> GenM p a) -> GenM p a)
-> (Gen RealWorld -> CategoricalProbability -> GenM p a)
-> GenM p a
forall a b. (a -> b) -> a -> b
$ [a]
-> Gen (PrimState (ReaderT (p, Gen RealWorld) IO))
-> CategoricalProbability
-> GenM p a
forall (m :: * -> *) a.
(Functor m, PrimMonad m) =>
[a] -> Gen (PrimState m) -> CategoricalProbability -> m a
sampleIn [a]
l

sampleFromString :: (p -> CategoricalProbability) -> B.ByteString -> GenM p Word8
sampleFromString :: forall p.
(p -> CategoricalProbability) -> ByteString -> GenM p Word8
sampleFromString p -> CategoricalProbability
p ByteString
s = (p -> CategoricalProbability)
-> (Gen RealWorld -> CategoricalProbability -> GenM p Word8)
-> GenM p Word8
forall p d x.
(p -> d) -> (Gen RealWorld -> d -> GenM p x) -> GenM p x
sampleWrapper p -> CategoricalProbability
p ((Gen RealWorld -> CategoricalProbability -> GenM p Word8)
 -> GenM p Word8)
-> (Gen RealWorld -> CategoricalProbability -> GenM p Word8)
-> GenM p Word8
forall a b. (a -> b) -> a -> b
$ ByteString
-> Gen (PrimState (ReaderT (p, Gen RealWorld) IO))
-> CategoricalProbability
-> GenM p Word8
forall (m :: * -> *).
(Functor m, PrimMonad m) =>
ByteString
-> Gen (PrimState m) -> CategoricalProbability -> m Word8
sampleInString ByteString
s

sampleBranch :: (p -> CategoricalProbability) -> [GenM p a] -> GenM p a
sampleBranch :: forall p a. (p -> CategoricalProbability) -> [GenM p a] -> GenM p a
sampleBranch p -> CategoricalProbability
p [GenM p a]
l = ReaderT (p, Gen RealWorld) IO (GenM p a) -> GenM p a
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (ReaderT (p, Gen RealWorld) IO (GenM p a) -> GenM p a)
-> ReaderT (p, Gen RealWorld) IO (GenM p a) -> GenM p a
forall a b. (a -> b) -> a -> b
$ (p -> CategoricalProbability)
-> [GenM p a] -> ReaderT (p, Gen RealWorld) IO (GenM p a)
forall p a. (p -> CategoricalProbability) -> [a] -> GenM p a
sampleFrom p -> CategoricalProbability
p [GenM p a]
l

sampleNum :: (p -> NumberProbability) -> GenM p Int
sampleNum :: forall p. (p -> NumberProbability) -> GenM p Int
sampleNum p -> NumberProbability
p = (p -> NumberProbability)
-> (Gen RealWorld -> NumberProbability -> GenM p Int) -> GenM p Int
forall p d x.
(p -> d) -> (Gen RealWorld -> d -> GenM p x) -> GenM p x
sampleWrapper p -> NumberProbability
p Gen RealWorld -> NumberProbability -> GenM p Int
Gen (PrimState (ReaderT (p, Gen RealWorld) IO))
-> NumberProbability -> GenM p Int
forall (m :: * -> *).
PrimMonad m =>
Gen (PrimState m) -> NumberProbability -> m Int
sampleNumberProbability

sampleN :: (p -> NumberProbability) -> GenM p b -> GenM p [b]
sampleN :: forall p b. (p -> NumberProbability) -> GenM p b -> GenM p [b]
sampleN p -> NumberProbability
p GenM p b
x = (p -> NumberProbability) -> GenM p Int
forall p. (p -> NumberProbability) -> GenM p Int
sampleNum p -> NumberProbability
p GenM p Int
-> (Int -> ReaderT (p, Gen RealWorld) IO [b])
-> ReaderT (p, Gen RealWorld) IO [b]
forall a b.
ReaderT (p, Gen RealWorld) IO a
-> (a -> ReaderT (p, Gen RealWorld) IO b)
-> ReaderT (p, Gen RealWorld) IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Int -> GenM p b -> ReaderT (p, Gen RealWorld) IO [b])
-> GenM p b -> Int -> ReaderT (p, Gen RealWorld) IO [b]
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> GenM p b -> ReaderT (p, Gen RealWorld) IO [b]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM GenM p b
x

sampleNE :: (p -> NumberProbability) -> GenM p b -> GenM p (NonEmpty b)
sampleNE :: forall p b.
(p -> NumberProbability) -> GenM p b -> GenM p (NonEmpty b)
sampleNE p -> NumberProbability
p GenM p b
x = (b -> [b] -> NonEmpty b)
-> GenM p b
-> ReaderT (p, Gen RealWorld) IO [b]
-> ReaderT (p, Gen RealWorld) IO (NonEmpty b)
forall a b c.
(a -> b -> c)
-> ReaderT (p, Gen RealWorld) IO a
-> ReaderT (p, Gen RealWorld) IO b
-> ReaderT (p, Gen RealWorld) IO c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 b -> [b] -> NonEmpty b
forall a. a -> [a] -> NonEmpty a
(:|) GenM p b
x (ReaderT (p, Gen RealWorld) IO [b]
 -> ReaderT (p, Gen RealWorld) IO (NonEmpty b))
-> ReaderT (p, Gen RealWorld) IO [b]
-> ReaderT (p, Gen RealWorld) IO (NonEmpty b)
forall a b. (a -> b) -> a -> b
$ (p -> NumberProbability)
-> GenM p b -> ReaderT (p, Gen RealWorld) IO [b]
forall p b. (p -> NumberProbability) -> GenM p b -> GenM p [b]
sampleN p -> NumberProbability
p GenM p b
x

sampleString ::
  (p -> NumberProbability) -> (p -> CategoricalProbability) -> B.ByteString -> GenM p B.ByteString
sampleString :: forall p.
(p -> NumberProbability)
-> (p -> CategoricalProbability) -> ByteString -> GenM p ByteString
sampleString p -> NumberProbability
np p -> CategoricalProbability
cp ByteString
s = [Word8] -> ByteString
B.pack ([Word8] -> ByteString)
-> ReaderT (p, Gen RealWorld) IO [Word8]
-> ReaderT (p, Gen RealWorld) IO ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (p -> NumberProbability)
-> GenM p Word8 -> ReaderT (p, Gen RealWorld) IO [Word8]
forall p b. (p -> NumberProbability) -> GenM p b -> GenM p [b]
sampleN p -> NumberProbability
np ((p -> CategoricalProbability) -> ByteString -> GenM p Word8
forall p.
(p -> CategoricalProbability) -> ByteString -> GenM p Word8
sampleFromString p -> CategoricalProbability
cp ByteString
s)

sampleNEString ::
  (p -> NumberProbability) -> (p -> CategoricalProbability) -> B.ByteString -> GenM p B.ByteString
sampleNEString :: forall p.
(p -> NumberProbability)
-> (p -> CategoricalProbability) -> ByteString -> GenM p ByteString
sampleNEString p -> NumberProbability
np p -> CategoricalProbability
cp ByteString
s = [Word8] -> ByteString
B.pack ([Word8] -> ByteString)
-> (NonEmpty Word8 -> [Word8]) -> NonEmpty Word8 -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty Word8 -> [Word8]
forall a. NonEmpty a -> [a]
toList (NonEmpty Word8 -> ByteString)
-> ReaderT (p, Gen RealWorld) IO (NonEmpty Word8)
-> ReaderT (p, Gen RealWorld) IO ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (p -> NumberProbability)
-> GenM p Word8 -> ReaderT (p, Gen RealWorld) IO (NonEmpty Word8)
forall p b.
(p -> NumberProbability) -> GenM p b -> GenM p (NonEmpty b)
sampleNE p -> NumberProbability
np ((p -> CategoricalProbability) -> ByteString -> GenM p Word8
forall p.
(p -> CategoricalProbability) -> ByteString -> GenM p Word8
sampleFromString p -> CategoricalProbability
cp ByteString
s)

deleteFirstOrdered :: Ord c => (a -> c) -> (b -> c) -> [a] -> [b] -> [a]
deleteFirstOrdered :: forall c a b. Ord c => (a -> c) -> (b -> c) -> [a] -> [b] -> [a]
deleteFirstOrdered a -> c
pa b -> c
pb [a]
la [b]
lb = case ([a]
la, [b]
lb) of
  (a
ha : [a]
ta, b
hb : [b]
tb) -> case c -> c -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (a -> c
pa a
ha) (b -> c
pb b
hb) of
    Ordering
LT -> a
ha a -> [a] -> [a]
forall a. a -> [a] -> [a]
: (a -> c) -> (b -> c) -> [a] -> [b] -> [a]
forall c a b. Ord c => (a -> c) -> (b -> c) -> [a] -> [b] -> [a]
deleteFirstOrdered a -> c
pa b -> c
pb [a]
ta [b]
lb
    Ordering
EQ -> (a -> c) -> (b -> c) -> [a] -> [b] -> [a]
forall c a b. Ord c => (a -> c) -> (b -> c) -> [a] -> [b] -> [a]
deleteFirstOrdered a -> c
pa b -> c
pb [a]
ta [b]
tb
    Ordering
GT -> (a -> c) -> (b -> c) -> [a] -> [b] -> [a]
forall c a b. Ord c => (a -> c) -> (b -> c) -> [a] -> [b] -> [a]
deleteFirstOrdered a -> c
pa b -> c
pb [a]
la [b]
tb
  ([a], [b])
_ -> [a]
la

merge :: Ord a => [a] -> [a] -> [a]
merge :: forall a. Ord a => [a] -> [a] -> [a]
merge [a]
la [a]
lb = case ([a]
la, [a]
lb) of
  (a
ha : [a]
ta, a
hb : [a]
tb) -> case a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare a
ha a
hb of
    Ordering
LT -> a
ha a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
forall a. Ord a => [a] -> [a] -> [a]
merge [a]
ta [a]
lb
    Ordering
EQ -> a
ha a -> [a] -> [a]
forall a. a -> [a] -> [a]
: a
hb a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
forall a. Ord a => [a] -> [a] -> [a]
merge [a]
ta [a]
tb
    Ordering
GT -> a
hb a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
forall a. Ord a => [a] -> [a] -> [a]
merge [a]
la [a]
tb
  ([a]
_, []) -> [a]
la
  ([a], [a])
_ -> [a]
lb

sampleFiltered :: (p -> CategoricalProbability) -> Int -> [Int] -> GenM p Int
sampleFiltered :: forall p.
(p -> CategoricalProbability) -> Int -> [Int] -> GenM p Int
sampleFiltered p -> CategoricalProbability
p Int
t [Int]
l = do
  Gen RealWorld
gen <- ((p, Gen RealWorld) -> Gen RealWorld)
-> ReaderT (p, Gen RealWorld) IO (Gen RealWorld)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (p, Gen RealWorld) -> Gen RealWorld
forall a b. (a, b) -> b
snd
  CategoricalProbability
d <- ((p, Gen RealWorld) -> CategoricalProbability)
-> ReaderT (p, Gen RealWorld) IO CategoricalProbability
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (((p, Gen RealWorld) -> CategoricalProbability)
 -> ReaderT (p, Gen RealWorld) IO CategoricalProbability)
-> ((p, Gen RealWorld) -> CategoricalProbability)
-> ReaderT (p, Gen RealWorld) IO CategoricalProbability
forall a b. (a -> b) -> a -> b
$ p -> CategoricalProbability
p (p -> CategoricalProbability)
-> ((p, Gen RealWorld) -> p)
-> (p, Gen RealWorld)
-> CategoricalProbability
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (p, Gen RealWorld) -> p
forall a b. (a, b) -> a
fst
  case CategoricalProbability
d of
    CPDiscrete NonEmpty Double
l ->
      [Int] -> Int -> Int
avoid [Int]
ll
        (Int -> Int) -> GenM p Int -> GenM p Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Prob (ReaderT (p, Gen RealWorld) IO) Int
-> Gen (PrimState (ReaderT (p, Gen RealWorld) IO)) -> GenM p Int
forall (m :: * -> *) a. Prob m a -> Gen (PrimState m) -> m a
sample ([(Double, Int)] -> Prob (ReaderT (p, Gen RealWorld) IO) Int
forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, PrimMonad m) =>
f (Double, a) -> Prob m a
discrete ([(Double, Int)] -> Prob (ReaderT (p, Gen RealWorld) IO) Int)
-> [(Double, Int)] -> Prob (ReaderT (p, Gen RealWorld) IO) Int
forall a b. (a -> b) -> a -> b
$ ((Double, Int) -> Int)
-> (Int -> Int) -> [(Double, Int)] -> [Int] -> [(Double, Int)]
forall c a b. Ord c => (a -> c) -> (b -> c) -> [a] -> [b] -> [a]
deleteFirstOrdered (Double, Int) -> Int
forall a b. (a, b) -> b
snd Int -> Int
forall a. a -> a
id ([Double] -> [Int] -> [(Double, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Int -> NonEmpty Double -> [Double]
forall a. Int -> NonEmpty a -> [a]
NE.take (Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) NonEmpty Double
l) [Int
0 .. Int
t]) [Int]
ll) Gen RealWorld
Gen (PrimState (ReaderT (p, Gen RealWorld) IO))
gen
    CPBiasedUniform [(Double, Int)]
l' Double
b ->
      let ll' :: [(Double, Int)]
ll' = ((Double, Int) -> Int)
-> (Int -> Int) -> [(Double, Int)] -> [Int] -> [(Double, Int)]
forall c a b. Ord c => (a -> c) -> (b -> c) -> [a] -> [b] -> [a]
deleteFirstOrdered (Double, Int) -> Int
forall a b. (a, b) -> b
snd Int -> Int
forall a. a -> a
id (Int -> [(Double, Int)] -> [(Double, Int)]
clean Int
t [(Double, Int)]
l') [Int]
ll
          uw :: Double
uw = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
- [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
ll Int -> Int -> Int
forall a. Num a => a -> a -> a
- [(Double, Int)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Double, Int)]
ll') Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b
       in Prob (ReaderT (p, Gen RealWorld) IO) (Maybe Int)
-> Gen (PrimState (ReaderT (p, Gen RealWorld) IO))
-> ReaderT (p, Gen RealWorld) IO (Maybe Int)
forall (m :: * -> *) a. Prob m a -> Gen (PrimState m) -> m a
sample ([(Double, Maybe Int)]
-> Prob (ReaderT (p, Gen RealWorld) IO) (Maybe Int)
forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, PrimMonad m) =>
f (Double, a) -> Prob m a
discrete ([(Double, Maybe Int)]
 -> Prob (ReaderT (p, Gen RealWorld) IO) (Maybe Int))
-> [(Double, Maybe Int)]
-> Prob (ReaderT (p, Gen RealWorld) IO) (Maybe Int)
forall a b. (a -> b) -> a -> b
$ (Double
uw, Maybe Int
forall a. Maybe a
Nothing) (Double, Maybe Int)
-> [(Double, Maybe Int)] -> [(Double, Maybe Int)]
forall a. a -> [a] -> [a]
: ((Double, Int) -> (Double, Maybe Int))
-> [(Double, Int)] -> [(Double, Maybe Int)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Double
x, Int
y) -> (Double
x, Int -> Maybe Int
forall a. a -> Maybe a
Just Int
y)) [(Double, Int)]
ll') Gen RealWorld
Gen (PrimState (ReaderT (p, Gen RealWorld) IO))
gen
            ReaderT (p, Gen RealWorld) IO (Maybe Int)
-> (Maybe Int -> GenM p Int) -> GenM p Int
forall a b.
ReaderT (p, Gen RealWorld) IO a
-> (a -> ReaderT (p, Gen RealWorld) IO b)
-> ReaderT (p, Gen RealWorld) IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= GenM p Int -> (Int -> GenM p Int) -> Maybe Int -> GenM p Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
              ( [Int] -> Int -> Int
avoid ([Int] -> [Int] -> [Int]
forall a. Ord a => [a] -> [a] -> [a]
merge [Int]
ll ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ ((Double, Int) -> Int) -> [(Double, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Double, Int) -> Int
forall a b. (a, b) -> b
snd [(Double, Int)]
ll')
                  (Int -> Int) -> GenM p Int -> GenM p Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Prob (ReaderT (p, Gen RealWorld) IO) Int
-> Gen (PrimState (ReaderT (p, Gen RealWorld) IO)) -> GenM p Int
forall (m :: * -> *) a. Prob m a -> Gen (PrimState m) -> m a
sample ((Int, Int) -> Prob (ReaderT (p, Gen RealWorld) IO) Int
forall (m :: * -> *) a.
(PrimMonad m, Variate a) =>
(a, a) -> Prob m a
uniformR (Int
0, Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
- [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
ll Int -> Int -> Int
forall a. Num a => a -> a -> a
- [(Double, Int)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Double, Int)]
ll')) Gen RealWorld
Gen (PrimState (ReaderT (p, Gen RealWorld) IO))
gen
              )
              Int -> GenM p Int
forall a. a -> ReaderT (p, Gen RealWorld) IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  where
    ll :: [Int]
ll = [Int] -> [Int]
forall a. Ord a => [a] -> [a]
sort ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
t) [Int]
l