module Data.Random.List where

import Data.Random.RVar
import Data.Random.Distribution.Uniform

import qualified System.Random.Shuffle as SRS
import Control.Monad

-- | A random variable returning an arbitrary element of the given list.
-- Every element has equal probability of being chosen.  Because it is a
-- pure 'RVar' it has no memory - that is, it \"draws with replacement.\"
randomElement :: [a] -> RVar a
randomElement :: forall a. [a] -> RVar a
randomElement = forall a (m :: * -> *). [a] -> RVarT m a
randomElementT


randomElementT :: [a] -> RVarT m a
randomElementT :: forall a (m :: * -> *). [a] -> RVarT m a
randomElementT [] = forall a. HasCallStack => [Char] -> a
error [Char]
"randomElementT: empty list!"
randomElementT [a]
xs = do
    Int
n <- forall a (m :: * -> *).
Distribution Uniform a =>
a -> a -> RVarT m a
uniformT Int
0 (forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs forall a. Num a => a -> a -> a
- Int
1)
    forall (m :: * -> *) a. Monad m => a -> m a
return ([a]
xs forall a. [a] -> Int -> a
!! Int
n)

-- | A random variable that returns the given list in an arbitrary shuffled
-- order.  Every ordering of the list has equal probability.
shuffle :: [a] -> RVar [a]
shuffle :: forall a. [a] -> RVar [a]
shuffle = forall a (m :: * -> *). [a] -> RVarT m [a]
shuffleT

shuffleT :: [a] -> RVarT m [a]
shuffleT :: forall a (m :: * -> *). [a] -> RVarT m [a]
shuffleT [] = forall (m :: * -> *) a. Monad m => a -> m a
return []
shuffleT [a]
xs = do
    [Int]
is <- forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (\a
_ Int
i -> forall a (m :: * -> *).
Distribution Uniform a =>
a -> a -> RVarT m a
uniformT Int
0 Int
i) (forall a. [a] -> [a]
tail [a]
xs) [Int
1..]

    forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. [a] -> [Int] -> [a]
SRS.shuffle [a]
xs (forall a. [a] -> [a]
reverse [Int]
is))

-- | A random variable that shuffles a list of a known length (or a list
-- prefix of the specified length). Useful for shuffling large lists when
-- the length is known in advance.  Avoids needing to traverse the list to
-- discover its length.  Each ordering has equal probability.
shuffleN :: Int -> [a] -> RVar [a]
shuffleN :: forall a. Int -> [a] -> RVar [a]
shuffleN = forall a (m :: * -> *). Int -> [a] -> RVarT m [a]
shuffleNT

shuffleNT :: Int -> [a] -> RVarT m [a]
shuffleNT :: forall a (m :: * -> *). Int -> [a] -> RVarT m [a]
shuffleNT Int
n [a]
xs = forall a (m :: * -> *). Int -> Int -> [a] -> RVarT m [a]
shuffleNofMT Int
n Int
n [a]
xs

-- | A random variable that selects N arbitrary elements of a list of known length M.
shuffleNofM :: Int -> Int -> [a] -> RVar [a]
shuffleNofM :: forall a. Int -> Int -> [a] -> RVar [a]
shuffleNofM = forall a (m :: * -> *). Int -> Int -> [a] -> RVarT m [a]
shuffleNofMT

shuffleNofMT :: Int -> Int -> [a] -> RVarT m [a]
shuffleNofMT :: forall a (m :: * -> *). Int -> Int -> [a] -> RVarT m [a]
shuffleNofMT Int
0 Int
_ [a]
_ = forall (m :: * -> *) a. Monad m => a -> m a
return []
shuffleNofMT Int
n Int
m [a]
xs
    | Int
n forall a. Ord a => a -> a -> Bool
> Int
m     = forall a. HasCallStack => [Char] -> a
error [Char]
"shuffleNofMT: n > m"
    | Int
n forall a. Ord a => a -> a -> Bool
>= Int
0    = do
        [Int]
is <- forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [forall a (m :: * -> *).
Distribution Uniform a =>
a -> a -> RVarT m a
uniformT Int
0 Int
i | Int
i <- forall a. Int -> [a] -> [a]
take Int
n [Int
mforall a. Num a => a -> a -> a
-Int
1, Int
mforall a. Num a => a -> a -> a
-Int
2 ..Int
1]]
        forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. Int -> [a] -> [a]
take Int
n forall a b. (a -> b) -> a -> b
$ forall a. [a] -> [Int] -> [a]
SRS.shuffle (forall a. Int -> [a] -> [a]
take Int
m [a]
xs) [Int]
is)
shuffleNofMT Int
_ Int
_ [a]
_ = forall a. HasCallStack => [Char] -> a
error [Char]
"shuffleNofMT: negative length specified"