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 :: [a] -> RVar a
randomElement = [a] -> RVar a
forall a (m :: * -> *). [a] -> RVarT m a
randomElementT


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

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

    [a] -> RVarT m [a]
forall (m :: * -> *) a. Monad m => a -> m a
return ([a] -> [Int] -> [a]
forall a. [a] -> [Int] -> [a]
SRS.shuffle [a]
xs ([Int] -> [Int]
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 :: Int -> [a] -> RVar [a]
shuffleN = Int -> [a] -> RVar [a]
forall a (m :: * -> *). Int -> [a] -> RVarT m [a]
shuffleNT

shuffleNT :: Int -> [a] -> RVarT m [a]
shuffleNT :: Int -> [a] -> RVarT m [a]
shuffleNT Int
n [a]
xs = Int -> Int -> [a] -> RVarT m [a]
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 :: Int -> Int -> [a] -> RVar [a]
shuffleNofM = Int -> Int -> [a] -> RVar [a]
forall a (m :: * -> *). Int -> Int -> [a] -> RVarT m [a]
shuffleNofMT

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