module Control.Monad.Random.Extras
(
shuffle
, shuffleSeq
, sample
, sampleSeq
, choiceExtract
, choiceExtractSeq
, choice
, safeChoice
, iterativeChoice
, choiceSeq
, safeChoiceSeq
, choiceArray
, choices
, safeChoices
, choicesArray
)
where
import Control.Monad (liftM)
import Control.Monad.Random (MonadRandom, getRandomR, getRandomRs)
import System.Random (Random)
import Data.List (mapAccumL, foldl')
import Data.Maybe (fromJust)
import qualified Data.Sequence as Seq
import Data.Sequence ((><), ViewL((:<)))
import qualified Data.Array.IArray as Arr
import qualified Data.Array
import Data.Array.IArray ((!))
(.:) :: (c -> c') -> (a -> b -> c) -> (a -> b -> c')
(.:) = (.).(.)
extract :: [a] -> Int -> Maybe ([a], a)
extract s i | null r = Nothing
| otherwise = Just (a ++ c, b)
where (a, r) = splitAt i s
(b : c) = r
extractSeq :: Seq.Seq a -> Int -> Maybe (Seq.Seq a, a)
extractSeq s i | Seq.null r = Nothing
| otherwise = Just (a >< c, b)
where (a, r) = Seq.splitAt i s
(b :< c) = Seq.viewl r
getRandomR' :: (MonadRandom m, Random a) => a -> a -> m a
getRandomR' = curry getRandomR
getRandomRNums :: (MonadRandom m, Random a, Num a) => [a] -> m [a]
getRandomRNums = mapM (getRandomR' 0)
backsaw :: Int -> [Int]
backsaw n = [n 1, n 2 .. 0]
shuffle :: (MonadRandom m) => [a] -> m [a]
shuffle = shuffleSeq . Seq.fromList
shuffleSeq :: (MonadRandom m) => Seq.Seq a -> m [a]
shuffleSeq s = do
samples <- getRandomRNums . backsaw $ Seq.length s
return (shuffleSeq' s samples)
shuffleSeq' :: Seq.Seq a -> [Int] -> [a]
shuffleSeq' = snd .: mapAccumL (fromJust .: extractSeq)
sample :: (MonadRandom m) => Int -> [a] -> m [a]
sample m = sampleSeq m . Seq.fromList
sampleSeq :: (MonadRandom m) => Int -> Seq.Seq a -> m [a]
sampleSeq m s = do
samples <- getRandomRNums . take m . backsaw $ Seq.length s
return (shuffleSeq' s samples)
choiceExtract :: (MonadRandom m) => [a] -> m (Maybe ([a], a))
choiceExtract [] = return Nothing
choiceExtract xs = extract xs `liftM` getRandomR (0, length xs 1)
choiceExtractSeq :: (MonadRandom m) => Seq.Seq a -> m (Maybe (Seq.Seq a, a))
choiceExtractSeq s | Seq.null s = return Nothing
| otherwise = extractSeq s `liftM` getRandomR (0, Seq.length s 1)
choice :: (MonadRandom m) => [a] -> m a
choice [] = error "Control.Monad.Random.Extras.choice: empty list"
choice xs = (xs !!) `liftM` getRandomR (0, length xs 1)
safeChoice :: (MonadRandom m) => [a] -> Maybe (m a)
safeChoice [] = Nothing
safeChoice xs = Just $ choice xs
iterativeChoice :: MonadRandom m => [a] -> m a
iterativeChoice xs = fst `liftM` foldl' stepM (return start) xs
where stepM x y = x >>= step y
step offered (old, n) = do
i <- getRandomR (0, n)
let new | i == 0 = offered
| otherwise = old
return $! new `seq` (new, n + 1)
start = (err, 0 :: Int)
err = error "Control.Monad.Random.Extras.iterativeChoice: empty list"
choiceSeq :: (MonadRandom m) => Seq.Seq a -> m a
choiceSeq s | Seq.null s = error "Control.Monad.Random.Extras.choiceSeq: empty sequence"
| otherwise = Seq.index s `liftM` getRandomR (0, Seq.length s 1)
safeChoiceSeq :: (MonadRandom m) => Seq.Seq a -> Maybe (m a)
safeChoiceSeq s | Seq.null s = Nothing
| otherwise = Just $ choiceSeq s
choiceArray :: (MonadRandom m, Arr.IArray arr a, Arr.Ix i, Random i) => arr i a -> m a
choiceArray v = (v !) `liftM` getRandomR (Arr.bounds v)
choices :: (MonadRandom m) => [a] -> m [a]
choices [] = error "Control.Monad.Random.Extras.choices: empty list"
choices xs = choicesArray $ Data.Array.listArray (1, length xs) xs
safeChoices :: (MonadRandom m) => [a] -> Maybe (m [a])
safeChoices [] = Nothing
safeChoices xs = Just $ choices xs
choicesArray :: (MonadRandom m, Arr.IArray arr a, Arr.Ix i, Random i) => arr i a -> m [a]
choicesArray v = map (v !) `liftM` getRandomRs (Arr.bounds v)