-- Copyright : (c) 2009 Oleg Kiselyov, Manlio Perillo

-- License : BSD3 (see LICENSE file)

--- http://okmij.org/ftp/Haskell/perfect-shuffle.txt

+-- <http://okmij.org/ftp/Haskell/perfect-shuffle.txt>

module System.Random.Shuffle

import System.Random (RandomGen, randomR)

--- A complete binary tree, of leaves and internal nodes.

+-- | A complete binary tree, of leaves and internal nodes.

-- Internal node: Node card l r

-- where card is the number of leaves under the node.

-- Invariant: card >=2. All internal tree nodes are always full.

--- Convert a sequence (e1...en) to a complete binary tree

+-- | Convert a sequence (e1...en) to a complete binary tree

buildTree :: [a] -> Tree a

buildTree = (fix growLevel) . (map Leaf)

join l@(Node ctl _ _) r@(Node ctr _ _) = Node (ctl + ctr) l r

--- Given a sequence (e1,...en) to shuffle, and a sequence

+-- | Given a sequence (e1,...en) to shuffle, and a sequence

-- (r1,...r[n-1]) of numbers such that r[i] is an independent sample

-- from a uniform random distribution [0..n-i], compute the

-- corresponding permutation of the input sequence.

| otherwise = let (e, r') = extractTree (n - cl) r

in (e, Node (c - 1) l r')

--- Given a sequence (e1,...en) to shuffle, its length, and a random

+-- | Given a sequence (e1,...en) to shuffle, its length, and a random

-- generator, compute the corresponding permutation of the input

shuffle' :: RandomGen gen => [a] -> Int -> gen -> [a]

-shuffle' elements len = shuffle elements . rseq len

+shuffle' elements len = shuffle elements . fst . rseq len

+-- | Given a sequence (e1,...en) to shuffle, its length, and a random

+-- generator, compute the corresponding permutation of the input

+-- sequence, return the permutation and the new state of the

+randomShuffle :: RandomGen gen => [a] -> Int -> gen -> ([a], gen)

+randomShuffle elements len g =

+ let (rs, g') = rseq len g

+ in (shuffle elements rs, g')

+-- | The sequence (r1,...r[n-1]) of numbers such that r[i] is an

+-- independent sample from a uniform random distribution

+rseq :: RandomGen gen => Int -> gen -> ([Int], gen)

+rseq n g = second lastGen . unzip $ rseq' (n - 1) g

- -- The sequence (r1,...r[n-1]) of numbers such that r[i] is an

- -- independent sample from a uniform random distribution

- rseq :: RandomGen gen => Int -> gen -> [Int]

- rseq n = fst . unzip . rseq' (n - 1)

- rseq' :: RandomGen gen => Int -> gen -> [(Int, gen)]

- rseq' i gen = (j, gen) : rseq' (i - 1) gen'

- (j, gen') = randomR (0, i) gen

+ rseq' :: RandomGen gen => Int -> gen -> [(Int, gen)]

+ | otherwise = let (j, gen') = randomR (0, i) gen

+ in (j, gen') : rseq' (i - 1) gen'

+ -- apply a function on the second element of a pair

+ second :: (b -> c) -> (a, b) -> (a, c)

+ second f (x,y) = (x, f y)

+ -- the last returned random number generator

+ lastGen [] = g -- didn't use the generator yet

+ lastGen gens = lastGen (drop 1 gens)