module RandomizedMonad (
	Randomized,
	runRandom, runRandomStd, runRandomNewStd,
	mrandomR, mrandom,
	withProb,
	filterRandomized
) where
import System.Random

-- Needs -XRank2Types
newtype Randomized a = Randomized (forall g. RandomGen g => (g -> a))

-- This implementation splits the RandomGen over and over.
-- It would also be possible to serialize everything and use a single RandomGen.
instance Monad Randomized where
	ma >>= amb = Randomized (\g -> let
			(g1, g2) = split g
			Randomized fa = ma
			a = fa g1
			Randomized fb = amb a
			in fb g2
		)
	return x = Randomized (const x)

runRandom :: RandomGen g => g -> Randomized a -> a
runRandom g (Randomized fa) = fa g

-- Conveniences
runRandomStd :: Randomized a -> IO a
runRandomStd ra = do
	g <- getStdGen
	return $ runRandom g ra

runRandomNewStd :: Randomized a -> IO a
runRandomNewStd ra = do
	g <- newStdGen
	return $ runRandom g ra

-- Monadic versions of random and randomR (to generate primitive-ish values)
mrandomR :: Random a => (a, a) -> Randomized a
mrandomR lohi = Randomized (\g -> fst (randomR lohi g))
mrandom :: Random a => Randomized a
mrandom = Randomized (\g -> fst (random g))

chooseCase :: Double -> [(Double, a)] -> a -> a
chooseCase val ifCs elseR = case ifCs of
	[] -> elseR
	(cutoff, theR):ifCt -> if val < cutoff
		then theR
		else chooseCase (val - cutoff) ifCt elseR

withProb :: [(Double, Randomized a)] -> Randomized a -> Randomized a
withProb ifCs elseR = do
	val <- mrandom
	chooseCase val ifCs elseR

-- Keep trying until we get what we want.
filterRandomized :: (a -> Bool) -> Randomized a -> Randomized a
filterRandomized f ra = do
	a <- ra
	if f a then return a else filterRandomized f ra