module RandomizedMonad (
	Randomized,
	msplit,
	runRandom1, runRandom, runRandomStd, runRandomNewStd,
	mrandomR, mrandom,
	withProb, withWeight,
	filterRandomized,
	indReplicateRandom, indRepeatRandom, indRandomArray
) where
import System.Random
import Data.Array.IArray
import Data.Ix

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

-- This implementation threads a single RandomGen through the whole process in
-- order to satisfy the monad laws.
instance Monad Randomized where
	ma >>= amb = Randomized (\g -> let
			Randomized fa = ma
			(a, g2) = fa g
			Randomized fb = amb a
			in fb g2
		)
	return x = Randomized (\g -> (x, g))

-- Splits the generator and runs the argument on the left generator while
-- threading the right generator on.  C.f. unsaveInterleaveIO.  Use this to
-- make a sub-calculation parallelizable and evolvable without breaking
-- same-seed reproducibility of the whole calculation.
msplit :: Randomized a -> Randomized a
msplit (Randomized fa) = Randomized
	(\g -> let (g1, g2) = split g in (fst (fa g1), g2))

runRandom1 :: RandomGen g => g -> Randomized a -> (a, g)
runRandom1 g (Randomized fa) = fa g

runRandom :: RandomGen g => g -> Randomized a -> a
runRandom g (Randomized fa) = fst (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)
mrandom :: Random a => Randomized a
mrandom = Randomized random
mrandomR :: Random a => (a, a) -> Randomized a
-- Eta-expand this one to keep GHC 6.6.1 on birdy happy.
mrandomR lohi = Randomized (\g -> randomR lohi 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

-- An if-elsif...else-style construct where each "if" has a probability.
withProb :: [(Double, Randomized a)] -> Randomized a -> Randomized a
withProb ifCs elseR = do
	val <- mrandom
	chooseCase val ifCs elseR

-- Like withProb, but without an else case and with the "probabilities" scaled
-- so that they sum to 1.
withWeight :: [(Double, Randomized a)] -> Randomized a
withWeight ifCs = do
	val <- mrandomR (0, sum (map fst ifCs))
	chooseCase val (tail ifCs) (snd (head ifCs))

-- 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

-- A randomized list of elements chosen independently from a distribution.
-- Each element is under msplit for parallelizability.
indReplicateRandom :: Int -> Randomized a -> Randomized [a]
indReplicateRandom n ra = sequence $ replicate n $ msplit ra

-- An infinite randomized list of elements chosen independently from a
-- distribution.  The list is under msplit to avoid an infinite loop when it is
-- bound.
indRepeatRandom :: Randomized a -> Randomized [a]
indRepeatRandom ra = msplit $ sequence $ repeat $ msplit ra

-- Produces an array of elements chosen independently from a distribution.
indRandomArray :: (IArray a e, Ix i) =>
	(i, i) -> Randomized e -> Randomized (a i e)
indRandomArray bds re = do
	list <- indReplicateRandom (rangeSize bds) re
	return (listArray bds list)