-- |
-- Module      :  ELynx.Tools.List
-- Description :  Additional tools for lists
-- Copyright   :  (c) Dominik Schrempf 2021
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Thu May  2 18:57:39 2019.
module ELynx.Tools.List
  ( -- * Lists
    sortListWithIndices,
    randomInsertList,
    shuffle,
    shuffleN,
    grabble,
  )
where

import Control.Monad
import Control.Monad.Primitive
import Control.Monad.ST
import Data.Function
import Data.List
import Data.Vector (Vector)
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as M
import System.Random.MWC

-- | Sort a list and also return original indices.
sortListWithIndices :: Ord a => [a] -> [(a, Int)]
sortListWithIndices :: [a] -> [(a, Int)]
sortListWithIndices [a]
xs = ((a, Int) -> (a, Int) -> Ordering) -> [(a, Int)] -> [(a, Int)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (a -> a -> Ordering)
-> ((a, Int) -> a) -> (a, Int) -> (a, Int) -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (a, Int) -> a
forall a b. (a, b) -> a
fst) ([(a, Int)] -> [(a, Int)]) -> [(a, Int)] -> [(a, Int)]
forall a b. (a -> b) -> a -> b
$ [a] -> [Int] -> [(a, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [a]
xs ([Int
0 ..] :: [Int])

-- | Insert element into random position of list.
randomInsertList :: PrimMonad m => a -> [a] -> Gen (PrimState m) -> m [a]
randomInsertList :: a -> [a] -> Gen (PrimState m) -> m [a]
randomInsertList a
e [a]
v Gen (PrimState m)
g = do
  let l :: Int
l = [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
v
  Int
i <- (Int, Int) -> Gen (PrimState m) -> m Int
forall a (m :: * -> *).
(Variate a, PrimMonad m) =>
(a, a) -> Gen (PrimState m) -> m a
uniformR (Int
0, Int
l) Gen (PrimState m)
g
  [a] -> m [a]
forall (m :: * -> *) a. Monad m => a -> m a
return ([a] -> m [a]) -> [a] -> m [a]
forall a b. (a -> b) -> a -> b
$ Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
i [a]
v [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a
e] [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
i [a]
v

-- | Shuffle a list.
shuffle :: PrimMonad m => [a] -> Gen (PrimState m) -> m [a]
shuffle :: [a] -> Gen (PrimState m) -> m [a]
shuffle [a]
xs Gen (PrimState m)
g = [[a]] -> [a]
forall a. [a] -> a
head ([[a]] -> [a]) -> m [[a]] -> m [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [a] -> Int -> Int -> Gen (PrimState m) -> m [[a]]
forall (m :: * -> *) a.
PrimMonad m =>
[a] -> Int -> Int -> Gen (PrimState m) -> m [[a]]
grabble [a]
xs Int
1 ([a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs) Gen (PrimState m)
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 (m :: * -> *) a.
PrimMonad m =>
[a] -> Int -> Int -> Gen (PrimState m) -> m [[a]]
grabble [a]
xs Int
n ([a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs)

-- | @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 :: PrimMonad m => [a] -> Int -> Int -> Gen (PrimState m) -> m [[a]]
grabble :: [a] -> Int -> Int -> Gen (PrimState m) -> m [[a]]
grabble [a]
xs Int
m Int
n Gen (PrimState m)
gen = do
  [[(Int, Int)]]
swapss <- Int -> m [(Int, Int)] -> m [[(Int, Int)]]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
m (m [(Int, Int)] -> m [[(Int, Int)]])
-> m [(Int, Int)] -> m [[(Int, Int)]]
forall a b. (a -> b) -> a -> b
$
    [Int] -> (Int -> m (Int, Int)) -> m [(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 -> m (Int, Int)) -> m [(Int, Int)])
-> (Int -> m (Int, Int)) -> m [(Int, Int)]
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
      Int
j <- (Int, Int) -> Gen (PrimState m) -> m Int
forall a (m :: * -> *).
(Variate a, PrimMonad m) =>
(a, a) -> Gen (PrimState m) -> m a
uniformR (Int
i, Int
l) Gen (PrimState m)
gen
      (Int, Int) -> m (Int, Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
i, Int
j)
  [[a]] -> m [[a]]
forall (m :: * -> *) a. Monad m => a -> m a
return ([[a]] -> m [[a]]) -> [[a]] -> m [[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