-- |
-- Module      :  Mcmc.Internal.Shuffle
-- Description :  Shuffle a list
-- Copyright   :  (c) Dominik Schrempf 2020
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Wed May 20 14:37:09 2020.
--
-- From https://wiki.haskell.org/Random_shuffle.
module Mcmc.Internal.Shuffle
  ( shuffle,
    shuffleN,
    grabble,
  )
where

import Control.Monad
import Control.Monad.ST
import Data.Vector (Vector)
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as M
import System.Random.MWC
  ( GenIO,
    uniformR,
  )

-- | Shuffle a list.
shuffle :: [a] -> GenIO -> IO [a]
shuffle :: [a] -> GenIO -> IO [a]
shuffle [a]
xs GenIO
g = [[a]] -> [a]
forall a. [a] -> a
head ([[a]] -> [a]) -> IO [[a]] -> IO [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [a] -> Int -> Int -> GenIO -> IO [[a]]
forall a. [a] -> Int -> Int -> GenIO -> IO [[a]]
grabble [a]
xs Int
1 ([a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs) GenIO
g

-- | Shuffle a list @n@ times.
shuffleN :: [a] -> Int -> GenIO -> IO [[a]]
shuffleN :: [a] -> Int -> GenIO -> IO [[a]]
shuffleN [a]
xs Int
n = [a] -> Int -> Int -> GenIO -> IO [[a]]
forall a. [a] -> Int -> Int -> GenIO -> IO [[a]]
grabble [a]
xs Int
n ([a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs)

-- -- Using System.Random.Shuffle. Speed is the same, so stay without additional dependency.
-- -- | Shuffle a list @n@ times.
-- shuffleN :: [a] -> Int -> GenIO -> IO [[a]]
-- shuffleN xs n g = replicateM n $ fmap (shuffle xs) (rseqM (length xs - 1) g)
--   where
--     rseqM :: Int -> GenIO -> IO [Int]
--     rseqM 0 _ = return []
--     rseqM i gen = liftM2 (:) (uniformR (0, i) gen) (rseqM (i - 1) gen)

-- | @grabble xs m n@ is /O(m*n')/, where @n' = min n (length xs)@. Choose @n'@
-- elements from @xs@, without replacement, and that @m@ times.
grabble :: [a] -> Int -> Int -> GenIO -> IO [[a]]
grabble :: [a] -> Int -> Int -> GenIO -> IO [[a]]
grabble [a]
xs Int
m Int
n GenIO
gen = do
  [[(Int, Int)]]
swapss <- Int -> IO [(Int, Int)] -> IO [[(Int, Int)]]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
m (IO [(Int, Int)] -> IO [[(Int, Int)]])
-> IO [(Int, Int)] -> IO [[(Int, Int)]]
forall a b. (a -> b) -> a -> b
$
    [Int] -> (Int -> IO (Int, Int)) -> IO [(Int, Int)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Int
0 .. Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
n] ((Int -> IO (Int, Int)) -> IO [(Int, Int)])
-> (Int -> IO (Int, Int)) -> IO [(Int, Int)]
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
      Int
j <- (Int, Int) -> GenIO -> IO Int
forall a (m :: * -> *).
(Variate a, PrimMonad m) =>
(a, a) -> Gen (PrimState m) -> m a
uniformR (Int
i, Int
l) GenIO
gen
      (Int, Int) -> IO (Int, Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
i, Int
j)
  [[a]] -> IO [[a]]
forall (m :: * -> *) a. Monad m => a -> m a
return ([[a]] -> IO [[a]]) -> [[a]] -> IO [[a]]
forall a b. (a -> b) -> a -> b
$ ([(Int, Int)] -> [a]) -> [[(Int, Int)]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (Vector a -> [a]
forall a. Vector a -> [a]
V.toList (Vector a -> [a])
-> ([(Int, Int)] -> Vector a) -> [(Int, Int)] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Vector a -> Vector a
forall a. Int -> Vector a -> Vector a
V.take Int
n (Vector a -> Vector a)
-> ([(Int, Int)] -> Vector a) -> [(Int, Int)] -> Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> [(Int, Int)] -> Vector a
forall a. Vector a -> [(Int, Int)] -> Vector a
swapElems ([a] -> Vector a
forall a. [a] -> Vector a
V.fromList [a]
xs)) [[(Int, Int)]]
swapss
  where
    l :: Int
l = [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

swapElems :: Vector a -> [(Int, Int)] -> Vector a
swapElems :: Vector a -> [(Int, Int)] -> Vector a
swapElems Vector a
xs [(Int, Int)]
swaps = (forall s. ST s (Vector a)) -> Vector a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector a)) -> Vector a)
-> (forall s. ST s (Vector a)) -> Vector a
forall a b. (a -> b) -> a -> b
$ do
  MVector s a
mxs <- Vector a -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Vector a -> m (MVector (PrimState m) a)
V.unsafeThaw Vector a
xs
  ((Int, Int) -> ST s ()) -> [(Int, Int)] -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((Int -> Int -> ST s ()) -> (Int, Int) -> ST s ()
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((Int -> Int -> ST s ()) -> (Int, Int) -> ST s ())
-> (Int -> Int -> ST s ()) -> (Int, Int) -> ST s ()
forall a b. (a -> b) -> a -> b
$ MVector (PrimState (ST s)) a -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> Int -> m ()
M.unsafeSwap MVector s a
MVector (PrimState (ST s)) a
mxs) [(Int, Int)]
swaps
  MVector (PrimState (ST s)) a -> ST s (Vector a)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
V.unsafeFreeze MVector s a
MVector (PrimState (ST s)) a
mxs