{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeApplications #-}

module Data.Vector.Shuffle
    ( -- * Simple
      shuffle

      -- * Advanced
    , mkSeed
    , shuffleWith
    ) where

import Prelude

import Control.Monad
    ( forM_ )
import Control.Monad.Trans.Class
    ( lift )
import Control.Monad.Trans.State.Strict
    ( evalStateT, state )
import Crypto.Hash
    ( hash )
import Crypto.Hash.Algorithms
    ( MD5 )
import Data.Text
    ( Text )
import Data.Vector.Mutable
    ( IOVector )
import Data.Word
    ( Word8 )
import System.Random
    ( RandomGen, newStdGen, randomR )

import qualified Data.ByteArray as BA
import qualified Data.ByteString as BS
import qualified Data.Text.Encoding as T
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV


-- | Generate a random generator seed from a text string
mkSeed :: Text -> Int
mkSeed :: Text -> Int
mkSeed = ByteString -> Int
toInt (ByteString -> Int) -> (Text -> ByteString) -> Text -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
quickHash (ByteString -> ByteString)
-> (Text -> ByteString) -> Text -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
T.encodeUtf16LE
  where
    quickHash :: ByteString -> ByteString
quickHash = Digest MD5 -> ByteString
forall bin bout.
(ByteArrayAccess bin, ByteArray bout) =>
bin -> bout
BA.convert (Digest MD5 -> ByteString)
-> (ByteString -> Digest MD5) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteArrayAccess ByteString, HashAlgorithm MD5) =>
ByteString -> Digest MD5
forall ba a.
(ByteArrayAccess ba, HashAlgorithm a) =>
ba -> Digest a
hash @_ @MD5
    toInt :: ByteString -> Int
toInt = (Int, Int) -> Int
forall a b. (a, b) -> b
snd ((Int, Int) -> Int)
-> (ByteString -> (Int, Int)) -> ByteString -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, Int) -> Word8 -> (Int, Int))
-> (Int, Int) -> ByteString -> (Int, Int)
forall a. (a -> Word8 -> a) -> a -> ByteString -> a
BS.foldl' (Int, Int) -> Word8 -> (Int, Int)
exponentiation (Int
0,Int
0)
      where
        exponentiation :: (Int, Int) -> Word8 -> (Int, Int)
        exponentiation :: (Int, Int) -> Word8 -> (Int, Int)
exponentiation (Int
e, Int
n) Word8
i = (Int
eInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1, Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
iInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
2Int -> Int -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^Int
e)

-- | Shuffles a list of elements.
--
-- >>> shuffle (outputs coinSel)
-- [...]
shuffle :: [a] -> IO [a]
shuffle :: [a] -> IO [a]
shuffle [a]
xs = IO StdGen
forall (m :: * -> *). MonadIO m => m StdGen
newStdGen IO StdGen -> (StdGen -> IO [a]) -> IO [a]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (StdGen -> [a] -> IO [a]) -> [a] -> StdGen -> IO [a]
forall a b c. (a -> b -> c) -> b -> a -> c
flip StdGen -> [a] -> IO [a]
forall g a. RandomGen g => g -> [a] -> IO [a]
shuffleWith [a]
xs

-- | Like 'shuffle', but from a given seed. 'shuffle' will use a randomly
-- generate seed using 'newStdGen' from @System.Random@.
--
-- __Properties:__
--
-- - @shuffleWith g es == shuffleWith g es@
-- - @∃Δ> 1. g ≠g', length es > Δ⇒ shuffleWith g es ≠shuffleWith g' es@
shuffleWith :: RandomGen g => g -> [a] -> IO [a]
shuffleWith :: g -> [a] -> IO [a]
shuffleWith g
seed = (IOVector a -> IO ()) -> [a] -> IO [a]
forall a. (IOVector a -> IO ()) -> [a] -> IO [a]
modifyInPlace ((IOVector a -> IO ()) -> [a] -> IO [a])
-> (IOVector a -> IO ()) -> [a] -> IO [a]
forall a b. (a -> b) -> a -> b
$ \IOVector a
v -> (StateT g IO () -> g -> IO ()) -> g -> StateT g IO () -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip StateT g IO () -> g -> IO ()
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT g
seed (StateT g IO () -> IO ()) -> StateT g IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    let (Int
lo, Int
hi) = (Int
0, IOVector a -> Int
forall s a. MVector s a -> Int
MV.length IOVector a
v Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
    [Int] -> (Int -> StateT g IO ()) -> StateT g IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
lo .. Int
hi] ((Int -> StateT g IO ()) -> StateT g IO ())
-> (Int -> StateT g IO ()) -> StateT g IO ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
      Int
j <- Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer -> Int) -> StateT g IO Integer -> StateT g IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (g -> (Integer, g)) -> StateT g IO Integer
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state ((Integer, Integer) -> g -> (Integer, g)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
lo, Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
hi))
      IO () -> StateT g IO ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (IO () -> StateT g IO ()) -> IO () -> StateT g IO ()
forall a b. (a -> b) -> a -> b
$ IOVector a -> Int -> Int -> IO ()
forall a. IOVector a -> Int -> Int -> IO ()
swapElems IOVector a
v Int
i Int
j
  where
    swapElems :: IOVector a -> Int -> Int -> IO ()
    swapElems :: IOVector a -> Int -> Int -> IO ()
swapElems IOVector a
v Int
i Int
j = do
        a
x <- MVector (PrimState IO) a -> Int -> IO a
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
MV.read IOVector a
MVector (PrimState IO) a
v Int
i
        a
y <- MVector (PrimState IO) a -> Int -> IO a
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
MV.read IOVector a
MVector (PrimState IO) a
v Int
j
        MVector (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write IOVector a
MVector (PrimState IO) a
v Int
i a
y
        MVector (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write IOVector a
MVector (PrimState IO) a
v Int
j a
x

    modifyInPlace :: forall a. (IOVector a -> IO ()) -> [a] -> IO [a]
    modifyInPlace :: (IOVector a -> IO ()) -> [a] -> IO [a]
modifyInPlace IOVector a -> IO ()
f [a]
xs = do
        IOVector a
v' <- Vector a -> IO (MVector (PrimState IO) a)
forall (m :: * -> *) a.
PrimMonad m =>
Vector a -> m (MVector (PrimState m) a)
V.thaw (Vector a -> IO (MVector (PrimState IO) a))
-> Vector a -> IO (MVector (PrimState IO) a)
forall a b. (a -> b) -> a -> b
$ [a] -> Vector a
forall a. [a] -> Vector a
V.fromList [a]
xs
        IOVector a -> IO ()
f IOVector a
v'
        Vector a -> [a]
forall a. Vector a -> [a]
V.toList (Vector a -> [a]) -> IO (Vector a) -> IO [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) a -> IO (Vector a)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
V.freeze IOVector a
MVector (PrimState IO) a
v'