--------------------------------------------------------------------------------
-- |
-- Module      :  HGeometry.Permutation.Shuffle
-- Copyright   :  (C) Frank Staals
-- License     :  see the LICENSE file
-- Maintainer  :  Frank Staals
--
-- Implements Fishyer-Yates shuffle.
--
--------------------------------------------------------------------------------
module HGeometry.Permutation.Shuffle
  ( shuffle
  , shuffleSeqInOut
  -- * For experimentation purposes only
  , shuffleSeqInOutOrig
  , shuffleIntMap
  , shuffleSeq
  ) where

import           Control.Lens (singular,ix,(&),(%%~),bimap)
import           Data.Foldable
import qualified Data.List as List
import           Data.Maybe (fromMaybe)
import qualified Data.Vector.Generic as V
import qualified Data.Vector.Generic.Mutable as MV
import           System.Random
import qualified VectorBuilder.Builder as Builder
import qualified VectorBuilder.MVector as Builder

import           Data.Sequence ((|>),(<|),Seq(..))
import qualified Data.Sequence as Seq
-- import           HGeometry.Sequence.NonEmpty (ViewR1(..))

import qualified Data.IntMap.Strict as IntMap

-- import qualified Data.Sequence.Internal as Internal

--------------------------------------------------------------------------------

-- | Fisher–Yates shuffle, which shuffles a list/foldable uniformly at random.
--
-- running time: \(O(n)\).
shuffle      :: forall vector gen a f.
                (V.Vector vector a, Foldable f, RandomGen gen) => gen -> f a -> vector a
shuffle :: forall (vector :: * -> *) gen a (f :: * -> *).
(Vector vector a, Foldable f, RandomGen gen) =>
gen -> f a -> vector a
shuffle gen
gen0 = Builder a -> vector a
construct (Builder a -> vector a) -> (f a -> Builder a) -> f a -> vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f a -> Builder a
forall (foldable :: * -> *) element.
Foldable foldable =>
foldable element -> Builder element
Builder.foldable
  where
    construct :: Builder a -> vector a
construct Builder a
b = (forall s. ST s (Mutable vector s a)) -> vector a
forall (v :: * -> *) a.
Vector v a =>
(forall s. ST s (Mutable v s a)) -> v a
V.create ((forall s. ST s (Mutable vector s a)) -> vector a)
-> (forall s. ST s (Mutable vector s a)) -> vector a
forall a b. (a -> b) -> a -> b
$ do
                    v <- Builder a -> ST s (Mutable vector s a)
forall (vector :: * -> * -> *) element s.
MVector vector element =>
Builder element -> ST s (vector s element)
Builder.build Builder a
b
                    for_ swaps $ \(Int
i,Int
j) ->
                      Mutable vector (PrimState (ST s)) a -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> Int -> m ()
MV.swap Mutable vector s a
Mutable vector (PrimState (ST s)) a
v Int
i Int
j
                    pure v
      where
        swaps :: [(Int, Int)]
swaps = ((Int, gen) -> Maybe ((Int, Int), (Int, gen)))
-> (Int, gen) -> [(Int, Int)]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
List.unfoldr (Int, gen) -> Maybe ((Int, Int), (Int, gen))
forall {a} {c}.
(Ord a, Num a, Enum a, RandomGen c, UniformRange a) =>
(a, c) -> Maybe ((a, a), (a, c))
f (Int -> Int
forall a. Enum a => a -> a
pred (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Builder a -> Int
forall element. Builder element -> Int
Builder.size Builder a
b, gen
gen0)
        f :: (a, c) -> Maybe ((a, a), (a, c))
f (a
i,c
gen)
          | a
i a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
1     = Maybe ((a, a), (a, c))
forall a. Maybe a
Nothing
          | Bool
otherwise = ((a, a), (a, c)) -> Maybe ((a, a), (a, c))
forall a. a -> Maybe a
Just (((a, a), (a, c)) -> Maybe ((a, a), (a, c)))
-> ((a, c) -> ((a, a), (a, c))) -> (a, c) -> Maybe ((a, a), (a, c))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> (a, a)) -> (c -> (a, c)) -> (a, c) -> ((a, a), (a, c))
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (a
i,) (a -> a
forall a. Enum a => a -> a
pred a
i,) ((a, c) -> Maybe ((a, a), (a, c)))
-> (a, c) -> Maybe ((a, a), (a, c))
forall a b. (a -> b) -> a -> b
$ (a, a) -> c -> (a, c)
forall g a. (RandomGen g, UniformRange a) => (a, a) -> g -> (a, g)
uniformR (a
0,a
i) c
gen


--------------------------------------------------------------------------------
-- * "Pure" versions

-- below are a bunch of implementations that try to avoid using vectors/mutability.  If
-- you care about that the "best" version is the shuffleSeqInOut version. The rest is all
-- slower.

-- | "Inside-out" version of Fissher-Yates shuffle that returns a Seq.  see
-- https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle#The_%22inside-out%22_algorithm
-- for details.
--
-- O(n\log n)
shuffleSeqInOut      :: (RandomGen gen, Foldable f) => gen -> f a -> Seq.Seq a
shuffleSeqInOut :: forall gen (f :: * -> *) a.
(RandomGen gen, Foldable f) =>
gen -> f a -> Seq a
shuffleSeqInOut gen
gen0 = (\(Acc gen
_ Seq a
s) -> Seq a
s) (Acc gen (Seq a) -> Seq a)
-> (f a -> Acc gen (Seq a)) -> f a -> Seq a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Acc gen (Seq a) -> a -> Acc gen (Seq a))
-> Acc gen (Seq a) -> f a -> Acc gen (Seq a)
forall b a. (b -> a -> b) -> b -> f a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Acc gen (Seq a) -> a -> Acc gen (Seq a)
forall {gen} {a}.
RandomGen gen =>
Acc gen (Seq a) -> a -> Acc gen (Seq a)
step (gen -> Seq a -> Acc gen (Seq a)
forall gen s. gen -> s -> Acc gen s
Acc gen
gen0 Seq a
forall a. Monoid a => a
mempty)
  where
    -- | sets the value at position i to x, and retrieves its current value.
    setAndRetrieve :: Int -> a -> Seq a -> (a, Seq a)
setAndRetrieve Int
i a
x Seq a
s = case Seq a
s Seq a -> Int -> Maybe a
forall a. Seq a -> Int -> Maybe a
Seq.!? Int
i of
      Maybe a
Nothing -> (a
x,Seq a
s)
      Just a
y  -> (a
y,Int -> a -> Seq a -> Seq a
forall a. Int -> a -> Seq a -> Seq a
Seq.update Int
i a
x Seq a
s)
    step :: Acc gen (Seq a) -> a -> Acc gen (Seq a)
step (Acc gen
gen Seq a
s) a
x = let (Int
j,gen
gen') = (Int, Int) -> gen -> (Int, gen)
forall g a. (RandomGen g, UniformRange a) => (a, a) -> g -> (a, g)
uniformR (Int
0,Seq a -> Int
forall a. Seq a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Seq a
s) gen
gen
                             (a
y,Seq a
s')   = Int -> a -> Seq a -> (a, Seq a)
forall {a}. Int -> a -> Seq a -> (a, Seq a)
setAndRetrieve Int
j a
x Seq a
s
                         in gen -> Seq a -> Acc gen (Seq a)
forall gen s. gen -> s -> Acc gen s
Acc gen
gen' (Seq a
s' Seq a -> a -> Seq a
forall a. Seq a -> a -> Seq a
|> a
y)
    -- main idea: for every next element x at position i, we generate a random index j <=
    -- i and place x at position j, and store the element y that was at position j at the
    -- new position i

data Acc gen s = Acc !gen !s
  deriving (forall a b. (a -> b) -> Acc gen a -> Acc gen b)
-> (forall a b. a -> Acc gen b -> Acc gen a) -> Functor (Acc gen)
forall a b. a -> Acc gen b -> Acc gen a
forall a b. (a -> b) -> Acc gen a -> Acc gen b
forall gen a b. a -> Acc gen b -> Acc gen a
forall gen a b. (a -> b) -> Acc gen a -> Acc gen b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall gen a b. (a -> b) -> Acc gen a -> Acc gen b
fmap :: forall a b. (a -> b) -> Acc gen a -> Acc gen b
$c<$ :: forall gen a b. a -> Acc gen b -> Acc gen a
<$ :: forall a b. a -> Acc gen b -> Acc gen a
Functor



--------------------------------------------------------------------------------

-- | "Inside-out" version of Fissher-Yates shuffle that returns a Seq.  see
-- https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle#The_%22inside-out%22_algorithm
-- for details.
--
-- O(n\log n)
shuffleSeqInOutOrig      :: (RandomGen gen, Foldable f) => gen -> f a -> Seq.Seq a
shuffleSeqInOutOrig :: forall gen (f :: * -> *) a.
(RandomGen gen, Foldable f) =>
gen -> f a -> Seq a
shuffleSeqInOutOrig gen
gen0 = (\(Acc gen
_ Seq a
s) -> Seq a
s) (Acc gen (Seq a) -> Seq a)
-> (f a -> Acc gen (Seq a)) -> f a -> Seq a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Acc gen (Seq a) -> a -> Acc gen (Seq a))
-> Acc gen (Seq a) -> f a -> Acc gen (Seq a)
forall b a. (b -> a -> b) -> b -> f a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Acc gen (Seq a) -> a -> Acc gen (Seq a)
forall {gen} {a}.
RandomGen gen =>
Acc gen (Seq a) -> a -> Acc gen (Seq a)
step (gen -> Seq a -> Acc gen (Seq a)
forall gen s. gen -> s -> Acc gen s
Acc gen
gen0 Seq a
forall a. Monoid a => a
mempty)
  where
    -- | sets the value at position i to x, and retrieves its current value.
    setAndRetrieve :: Index t -> IxValue t -> t -> SP (IxValue t) t
setAndRetrieve Index t
i IxValue t
x t
s = t
st -> (t -> SP (IxValue t) t) -> SP (IxValue t) t
forall a b. a -> (a -> b) -> b
&Index t -> Traversal' t (IxValue t)
forall m. Ixed m => Index m -> Traversal' m (IxValue m)
ix Index t
i ((IxValue t -> SP (IxValue t) (IxValue t))
 -> t -> SP (IxValue t) t)
-> (IxValue t -> SP (IxValue t) (IxValue t))
-> t
-> SP (IxValue t) t
forall {k} (f :: k -> *) s (t :: k) a (b :: k).
LensLike f s t a b -> LensLike f s t a b
%%~ \IxValue t
y -> Maybe (IxValue t) -> IxValue t -> SP (IxValue t) (IxValue t)
forall a b. Maybe a -> b -> SP a b
SP (IxValue t -> Maybe (IxValue t)
forall a. a -> Maybe a
Just IxValue t
y) IxValue t
x
      -- the SP here is very important; if we use a lazy pair this about 4x lower
      -- same for the ! on y below here.

    step :: Acc gen (Seq a) -> a -> Acc gen (Seq a)
step (Acc gen
gen Seq a
s) a
x = let (!Int
j,gen
gen') = (Int, Int) -> gen -> (Int, gen)
forall g a. (RandomGen g, UniformRange a) => (a, a) -> g -> (a, g)
uniformR (Int
0,Seq a -> Int
forall a. Seq a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Seq a
s) gen
gen
                             SP Maybe (IxValue (Seq a))
my Seq a
s'  = Index (Seq a)
-> IxValue (Seq a) -> Seq a -> SP (IxValue (Seq a)) (Seq a)
forall {t}. Ixed t => Index t -> IxValue t -> t -> SP (IxValue t) t
setAndRetrieve Int
Index (Seq a)
j a
IxValue (Seq a)
x Seq a
s
                             !y :: a
y        = a -> Maybe a -> a
forall a. a -> Maybe a -> a
fromMaybe a
x Maybe a
Maybe (IxValue (Seq a))
my
                         in gen -> Seq a -> Acc gen (Seq a)
forall gen s. gen -> s -> Acc gen s
Acc gen
gen' (Seq a
s' Seq a -> a -> Seq a
forall a. Seq a -> a -> Seq a
|> a
y)
    -- main idea: for every next element x at position i, we generate a random index j <=
    -- i and place x at position j, and store the element y that was at position j at the
    -- new position i

--------------------------------------------------------------------------------

data SP a b = SP !(Maybe a) !b
  deriving ((forall a b. (a -> b) -> SP a a -> SP a b)
-> (forall a b. a -> SP a b -> SP a a) -> Functor (SP a)
forall a b. a -> SP a b -> SP a a
forall a b. (a -> b) -> SP a a -> SP a b
forall a a b. a -> SP a b -> SP a a
forall a a b. (a -> b) -> SP a a -> SP a b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a a b. (a -> b) -> SP a a -> SP a b
fmap :: forall a b. (a -> b) -> SP a a -> SP a b
$c<$ :: forall a a b. a -> SP a b -> SP a a
<$ :: forall a b. a -> SP a b -> SP a a
Functor)

instance Applicative (SP a) where
  pure :: forall a. a -> SP a a
pure a
x            = Maybe a -> a -> SP a a
forall a b. Maybe a -> b -> SP a b
SP Maybe a
forall a. Maybe a
Nothing a
x
  SP Maybe a
a a -> b
f <*> :: forall a b. SP a (a -> b) -> SP a a -> SP a b
<*> SP Maybe a
_ a
x = Maybe a -> b -> SP a b
forall a b. Maybe a -> b -> SP a b
SP Maybe a
a (a -> b
f a
x)



--------------------------------------------------------------------------------

-- | Returns a strict IntMap
shuffleIntMap      :: (RandomGen gen, Foldable f) => gen -> f a -> IntMap.IntMap a
shuffleIntMap :: forall gen (f :: * -> *) a.
(RandomGen gen, Foldable f) =>
gen -> f a -> IntMap a
shuffleIntMap gen
gen0 = IntMap a -> IntMap a
build (IntMap a -> IntMap a) -> (f a -> IntMap a) -> f a -> IntMap a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Int, a)] -> IntMap a
forall a. [(Int, a)] -> IntMap a
IntMap.fromList ([(Int, a)] -> IntMap a) -> (f a -> [(Int, a)]) -> f a -> IntMap a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] ([a] -> [(Int, a)]) -> (f a -> [a]) -> f a -> [(Int, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f a -> [a]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList
  where
    build :: IntMap a -> IntMap a
build IntMap a
m = (IntMap a -> (Int, Int) -> IntMap a)
-> IntMap a -> [(Int, Int)] -> IntMap a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' IntMap a -> (Int, Int) -> IntMap a
forall {a}. IntMap a -> (Int, Int) -> IntMap a
swap IntMap a
m ([(Int, Int)] -> IntMap a) -> [(Int, Int)] -> IntMap a
forall a b. (a -> b) -> a -> b
$ ((Int, gen) -> Maybe ((Int, Int), (Int, gen)))
-> (Int, gen) -> [(Int, Int)]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
List.unfoldr (Int, gen) -> Maybe ((Int, Int), (Int, gen))
forall {a} {c}.
(Ord a, Num a, Enum a, RandomGen c, UniformRange a) =>
(a, c) -> Maybe ((a, a), (a, c))
f (Int -> Int
forall a. Enum a => a -> a
pred (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ IntMap a -> Int
forall a. IntMap a -> Int
IntMap.size IntMap a
m, gen
gen0)
      where
        f :: (a, c) -> Maybe ((a, a), (a, c))
f (a
i,c
gen)
          | a
i a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
1     = Maybe ((a, a), (a, c))
forall a. Maybe a
Nothing
          | Bool
otherwise = ((a, a), (a, c)) -> Maybe ((a, a), (a, c))
forall a. a -> Maybe a
Just (((a, a), (a, c)) -> Maybe ((a, a), (a, c)))
-> ((a, c) -> ((a, a), (a, c))) -> (a, c) -> Maybe ((a, a), (a, c))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> (a, a)) -> (c -> (a, c)) -> (a, c) -> ((a, a), (a, c))
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (a
i,) (a -> a
forall a. Enum a => a -> a
pred a
i,) ((a, c) -> Maybe ((a, a), (a, c)))
-> (a, c) -> Maybe ((a, a), (a, c))
forall a b. (a -> b) -> a -> b
$ (a, a) -> c -> (a, c)
forall g a. (RandomGen g, UniformRange a) => (a, a) -> g -> (a, g)
uniformR (a
0,a
i) c
gen

    swap :: IntMap a -> (Int, Int) -> IntMap a
swap IntMap a
m (Int
i,Int
j) = Int -> a -> IntMap a -> IntMap a
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
i (IntMap a
m IntMap a -> Int -> a
forall a. IntMap a -> Int -> a
IntMap.! Int
j) (IntMap a -> IntMap a)
-> (IntMap a -> IntMap a) -> IntMap a -> IntMap a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> a -> IntMap a -> IntMap a
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
j (IntMap a
m IntMap a -> Int -> a
forall a. IntMap a -> Int -> a
IntMap.! Int
i) (IntMap a -> IntMap a) -> IntMap a -> IntMap a
forall a b. (a -> b) -> a -> b
$ IntMap a
m


-- | Version of Fissher-Yates shuffle that returns a Seq.
--
-- O(n\log n)
shuffleSeq      :: (RandomGen gen, Foldable f) => gen -> f a -> Seq.Seq a
shuffleSeq :: forall gen (f :: * -> *) a.
(RandomGen gen, Foldable f) =>
gen -> f a -> Seq a
shuffleSeq gen
gen0 = Seq a -> gen -> Seq a -> Seq a
forall {t} {a}. RandomGen t => Seq a -> t -> Seq a -> Seq a
build Seq a
forall a. Monoid a => a
mempty gen
gen0 (Seq a -> Seq a) -> (f a -> Seq a) -> f a -> Seq a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Seq a) -> f a -> Seq a
forall m a. Monoid m => (a -> m) -> f a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> Seq a
forall a. a -> Seq a
Seq.singleton
  where
    setAndRetrieve :: Index t -> IxValue t -> t -> Acc (IxValue t) t
setAndRetrieve Index t
i IxValue t
x t
s = t
st -> (t -> Acc (IxValue t) t) -> Acc (IxValue t) t
forall a b. a -> (a -> b) -> b
&Traversing (->) (Acc (IxValue t)) t t (IxValue t) (IxValue t)
-> (IxValue t -> Acc (IxValue t) (IxValue t))
-> t
-> Acc (IxValue t) t
forall (p :: * -> * -> *) (f :: * -> *) s t a.
(HasCallStack, Conjoined p, Functor f) =>
Traversing p f s t a a -> Over p f s t a a
singular (Index t -> Traversal' t (IxValue t)
forall m. Ixed m => Index m -> Traversal' m (IxValue m)
ix Index t
i) ((IxValue t -> Acc (IxValue t) (IxValue t))
 -> t -> Acc (IxValue t) t)
-> (IxValue t -> Acc (IxValue t) (IxValue t))
-> t
-> Acc (IxValue t) t
forall {k} (f :: k -> *) s (t :: k) a (b :: k).
LensLike f s t a b -> LensLike f s t a b
%%~ \IxValue t
y -> IxValue t -> IxValue t -> Acc (IxValue t) (IxValue t)
forall gen s. gen -> s -> Acc gen s
Acc IxValue t
y IxValue t
x
    build :: Seq a -> t -> Seq a -> Seq a
build Seq a
s t
gen = \case
      Seq a
Empty             -> Seq a
s
      (Seq a
remaining :|> a
x) -> let i :: Int
i              = Seq a -> Int
forall a. Seq a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Seq a
remaining
                               (Int
j,t
gen')       = (Int, Int) -> t -> (Int, t)
forall g a. (RandomGen g, UniformRange a) => (a, a) -> g -> (a, g)
uniformR (Int
0,Int
i) t
gen
                               (Acc IxValue (Seq a)
y Seq a
remaining')
                                 | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
j     = Index (Seq a)
-> IxValue (Seq a) -> Seq a -> Acc (IxValue (Seq a)) (Seq a)
forall {t}.
Ixed t =>
Index t -> IxValue t -> t -> Acc (IxValue t) t
setAndRetrieve Int
Index (Seq a)
j a
IxValue (Seq a)
x Seq a
remaining
                                 | Bool
otherwise  = a -> Seq a -> Acc a (Seq a)
forall gen s. gen -> s -> Acc gen s
Acc a
x Seq a
remaining
                           in Seq a -> t -> Seq a -> Seq a
build (a
IxValue (Seq a)
y a -> Seq a -> Seq a
forall a. a -> Seq a -> Seq a
<| Seq a
s) t
gen' Seq a
remaining'