module Data.Vector.Utils where

import Control.Monad
import Control.Monad.Primitive
import Data.Function
import qualified Data.Vector.Fusion.Stream.Monadic as MS
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM
import qualified Data.Vector.Unboxed as U

import My.Prelude

chunks :: (G.Vector u a, G.Vector v (u a)) => Int -> u a -> v (u a)
chunks :: forall (u :: * -> *) a (v :: * -> *).
(Vector u a, Vector v (u a)) =>
Int -> u a -> v (u a)
chunks Int
n = (u a -> Maybe (u a, u a)) -> u a -> v (u a)
forall (v :: * -> *) a b.
Vector v a =>
(b -> Maybe (a, b)) -> b -> v a
G.unfoldr ((u a -> Maybe (u a, u a)) -> u a -> v (u a))
-> (u a -> Maybe (u a, u a)) -> u a -> v (u a)
forall a b. (a -> b) -> a -> b
$ \u a
u ->
  if u a -> Bool
forall (v :: * -> *) a. Vector v a => v a -> Bool
G.null u a
u then Maybe (u a, u a)
forall a. Maybe a
Nothing else (u a, u a) -> Maybe (u a, u a)
forall a. a -> Maybe a
Just ((u a, u a) -> Maybe (u a, u a)) -> (u a, u a) -> Maybe (u a, u a)
forall a b. (a -> b) -> a -> b
$ Int -> u a -> (u a, u a)
forall (v :: * -> *) a. Vector v a => Int -> v a -> (v a, v a)
G.splitAt Int
n u a
u

group :: (Eq a, G.Vector u a, G.Vector v (u a)) => u a -> v (u a)
group :: forall a (u :: * -> *) (v :: * -> *).
(Eq a, Vector u a, Vector v (u a)) =>
u a -> v (u a)
group = (a -> a -> Bool) -> u a -> v (u a)
forall (u :: * -> *) a (v :: * -> *).
(Vector u a, Vector v (u a)) =>
(a -> a -> Bool) -> u a -> v (u a)
groupBy a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(==)

groupBy :: (G.Vector u a, G.Vector v (u a)) => (a -> a -> Bool) -> u a -> v (u a)
groupBy :: forall (u :: * -> *) a (v :: * -> *).
(Vector u a, Vector v (u a)) =>
(a -> a -> Bool) -> u a -> v (u a)
groupBy a -> a -> Bool
eq = (u a -> Maybe (u a, u a)) -> u a -> v (u a)
forall (v :: * -> *) a b.
Vector v a =>
(b -> Maybe (a, b)) -> b -> v a
G.unfoldr ((u a -> Maybe (u a, u a)) -> u a -> v (u a))
-> (u a -> Maybe (u a, u a)) -> u a -> v (u a)
forall a b. (a -> b) -> a -> b
$ \u a
u ->
  if u a -> Bool
forall (v :: * -> *) a. Vector v a => v a -> Bool
G.null u a
u
    then Maybe (u a, u a)
forall a. Maybe a
Nothing
    else case u a -> a
forall (v :: * -> *) a. Vector v a => v a -> a
G.unsafeHead u a
u of
      a
u0 -> (u a, u a) -> Maybe (u a, u a)
forall a. a -> Maybe a
Just ((u a, u a) -> Maybe (u a, u a)) -> (u a, u a) -> Maybe (u a, u a)
forall a b. (a -> b) -> a -> b
$ (a -> Bool) -> u a -> (u a, u a)
forall (v :: * -> *) a.
Vector v a =>
(a -> Bool) -> v a -> (v a, v a)
G.span (a -> a -> Bool
eq a
u0) u a
u

tuples2 :: (G.Vector v a, G.Vector v (a, a)) => v a -> v (a, a)
tuples2 :: forall (v :: * -> *) a.
(Vector v a, Vector v (a, a)) =>
v a -> v (a, a)
tuples2 = (v a -> Maybe ((a, a), v a)) -> v a -> v (a, a)
forall (v :: * -> *) a b.
Vector v a =>
(b -> Maybe (a, b)) -> b -> v a
G.unfoldr ((v a -> Maybe ((a, a), v a)) -> v a -> v (a, a))
-> (v a -> Maybe ((a, a), v a)) -> v a -> v (a, a)
forall a b. (a -> b) -> a -> b
$ \v a
v ->
  if v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
v Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2
    then
      ( let !x :: a
x = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
0
            !y :: a
y = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
1
         in ((a, a), v a) -> Maybe ((a, a), v a)
forall a. a -> Maybe a
Just ((a
x, a
y), Int -> v a -> v a
forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
G.drop Int
2 v a
v)
      )
    else Maybe ((a, a), v a)
forall a. Maybe a
Nothing

tuples2N :: (G.Vector v a, G.Vector v (a, a)) => Int -> v a -> v (a, a)
tuples2N :: forall (v :: * -> *) a.
(Vector v a, Vector v (a, a)) =>
Int -> v a -> v (a, a)
tuples2N Int
n = Int -> (v a -> ((a, a), v a)) -> v a -> v (a, a)
forall (v :: * -> *) a b.
Vector v a =>
Int -> (b -> (a, b)) -> b -> v a
G.unfoldrExactN Int
n ((v a -> ((a, a), v a)) -> v a -> v (a, a))
-> (v a -> ((a, a), v a)) -> v a -> v (a, a)
forall a b. (a -> b) -> a -> b
$ \v a
v ->
  let !x :: a
x = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
0
      !y :: a
y = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
1
   in ((a
x, a
y), Int -> v a -> v a
forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
G.drop Int
2 v a
v)

tuples3 :: (G.Vector v a, G.Vector v (a, a, a)) => v a -> v (a, a, a)
tuples3 :: forall (v :: * -> *) a.
(Vector v a, Vector v (a, a, a)) =>
v a -> v (a, a, a)
tuples3 = (v a -> Maybe ((a, a, a), v a)) -> v a -> v (a, a, a)
forall (v :: * -> *) a b.
Vector v a =>
(b -> Maybe (a, b)) -> b -> v a
G.unfoldr ((v a -> Maybe ((a, a, a), v a)) -> v a -> v (a, a, a))
-> (v a -> Maybe ((a, a, a), v a)) -> v a -> v (a, a, a)
forall a b. (a -> b) -> a -> b
$ \v a
v ->
  if v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
v Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
3
    then
      ( let !x :: a
x = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
0
            !y :: a
y = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
1
            !z :: a
z = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
2
         in ((a, a, a), v a) -> Maybe ((a, a, a), v a)
forall a. a -> Maybe a
Just ((a
x, a
y, a
z), Int -> v a -> v a
forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
G.drop Int
3 v a
v)
      )
    else Maybe ((a, a, a), v a)
forall a. Maybe a
Nothing

tuples3N :: (G.Vector v a, G.Vector v (a, a, a)) => Int -> v a -> v (a, a, a)
tuples3N :: forall (v :: * -> *) a.
(Vector v a, Vector v (a, a, a)) =>
Int -> v a -> v (a, a, a)
tuples3N Int
n = Int -> (v a -> ((a, a, a), v a)) -> v a -> v (a, a, a)
forall (v :: * -> *) a b.
Vector v a =>
Int -> (b -> (a, b)) -> b -> v a
G.unfoldrExactN Int
n ((v a -> ((a, a, a), v a)) -> v a -> v (a, a, a))
-> (v a -> ((a, a, a), v a)) -> v a -> v (a, a, a)
forall a b. (a -> b) -> a -> b
$ \v a
v ->
  let !x :: a
x = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
0
      !y :: a
y = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
1
      !z :: a
z = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
2
   in ((a
x, a
y, a
z), Int -> v a -> v a
forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
G.drop Int
3 v a
v)

tuples4 :: (G.Vector v a, G.Vector v (a, a, a, a)) => v a -> v (a, a, a, a)
tuples4 :: forall (v :: * -> *) a.
(Vector v a, Vector v (a, a, a, a)) =>
v a -> v (a, a, a, a)
tuples4 = (v a -> Maybe ((a, a, a, a), v a)) -> v a -> v (a, a, a, a)
forall (v :: * -> *) a b.
Vector v a =>
(b -> Maybe (a, b)) -> b -> v a
G.unfoldr ((v a -> Maybe ((a, a, a, a), v a)) -> v a -> v (a, a, a, a))
-> (v a -> Maybe ((a, a, a, a), v a)) -> v a -> v (a, a, a, a)
forall a b. (a -> b) -> a -> b
$ \v a
v ->
  if v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
v Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
3
    then
      ( let !x :: a
x = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
0
            !y :: a
y = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
1
            !z :: a
z = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
2
            !w :: a
w = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
3
         in ((a, a, a, a), v a) -> Maybe ((a, a, a, a), v a)
forall a. a -> Maybe a
Just ((a
x, a
y, a
z, a
w), Int -> v a -> v a
forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
G.drop Int
4 v a
v)
      )
    else Maybe ((a, a, a, a), v a)
forall a. Maybe a
Nothing

tuples4N ::
  (G.Vector v a, G.Vector v (a, a, a, a)) =>
  Int ->
  v a ->
  v (a, a, a, a)
tuples4N :: forall (v :: * -> *) a.
(Vector v a, Vector v (a, a, a, a)) =>
Int -> v a -> v (a, a, a, a)
tuples4N Int
n = Int -> (v a -> ((a, a, a, a), v a)) -> v a -> v (a, a, a, a)
forall (v :: * -> *) a b.
Vector v a =>
Int -> (b -> (a, b)) -> b -> v a
G.unfoldrExactN Int
n ((v a -> ((a, a, a, a), v a)) -> v a -> v (a, a, a, a))
-> (v a -> ((a, a, a, a), v a)) -> v a -> v (a, a, a, a)
forall a b. (a -> b) -> a -> b
$ \v a
v ->
  let !x :: a
x = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
0
      !y :: a
y = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
1
      !z :: a
z = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
2
      !w :: a
w = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
3
   in ((a
x, a
y, a
z, a
w), Int -> v a -> v a
forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
G.drop Int
4 v a
v)

{- |
>>> import qualified Data.Vector.Unboxed as U
>>> transpose 2 3 (U.fromList "abcdef")
"adbecf"
-}
transpose :: (G.Vector v a) => Int -> Int -> v a -> v a
transpose :: forall (v :: * -> *) a. Vector v a => Int -> Int -> v a -> v a
transpose Int
h Int
w v a
vec = (forall s. ST s (Mutable v s a)) -> v a
forall (v :: * -> *) a.
Vector v a =>
(forall s. ST s (Mutable v s a)) -> v a
G.create ((forall s. ST s (Mutable v s a)) -> v a)
-> (forall s. ST s (Mutable v s a)) -> v a
forall a b. (a -> b) -> a -> b
$ do
  Mutable v s a
buf <- Int -> ST s (Mutable v (PrimState (ST s)) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
GM.unsafeNew (Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
h)
  Int -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *). Monad m => Int -> (Int -> m ()) -> m ()
rep Int
h ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
x -> do
    (Int -> a -> ST s Int) -> Int -> v a -> ST s ()
forall (m :: * -> *) (v :: * -> *) b a.
(Monad m, Vector v b) =>
(a -> b -> m a) -> a -> v b -> m ()
G.foldM'_
      ( \Int
pos a
mxy -> do
          Mutable v (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
GM.write Mutable v s a
Mutable v (PrimState (ST s)) a
buf Int
pos a
mxy
          Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$! Int
pos Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
h
      )
      Int
x
      (Int -> Int -> v a -> v a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
Int -> Int -> v a -> v a
G.slice (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
w) Int
w v a
vec)
  Mutable v s a -> ST s (Mutable v s a)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return Mutable v s a
buf

{- |
>>> U.modify (void . nextPermutation) $ U.fromList "abc"
"acb"
>>> U.modify (void . nextPermutation) $ U.fromList "cba"
"cba"
>>> U.modify (void . nextPermutation) $ U.fromList "a"
"a"
>>> U.modify (void . nextPermutation) $ U.fromList ""
""
-}
nextPermutation :: (GM.MVector mv a, Ord a, PrimMonad m) => mv (PrimState m) a -> m Bool
nextPermutation :: forall (mv :: * -> * -> *) a (m :: * -> *).
(MVector mv a, Ord a, PrimMonad m) =>
mv (PrimState m) a -> m Bool
nextPermutation mv (PrimState m) a
v = do
  ((Int -> Int -> Int -> m Bool) -> Int -> Int -> Int -> m Bool)
-> Int -> Int -> Int -> m Bool
forall a. (a -> a) -> a
fix
    ( \Int -> Int -> Int -> m Bool
loop !Int
i !Int
k !Int
l ->
        if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n
          then do
            a
prev <- mv (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead mv (PrimState m) a
v (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
            a
cur <- mv (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead mv (PrimState m) a
v Int
i
            if a
prev a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
cur
              then Int -> Int -> Int -> m Bool
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
i
              else
                if Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
nothing
                  then do
                    a
vk <- mv (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead mv (PrimState m) a
v Int
k
                    if a
vk a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
cur
                      then Int -> Int -> Int -> m Bool
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
k Int
i
                      else Int -> Int -> Int -> m Bool
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
k Int
l
                  else Int -> Int -> Int -> m Bool
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
k Int
l
          else do
            if Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
nothing
              then do
                mv (PrimState m) a -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> Int -> m ()
GM.unsafeSwap mv (PrimState m) a
v Int
k Int
l
                mv (PrimState m) a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> m ()
GM.reverse (mv (PrimState m) a -> m ()) -> mv (PrimState m) a -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> mv (PrimState m) a -> mv (PrimState m) a
forall (v :: * -> * -> *) a s. MVector v a => Int -> v s a -> v s a
GM.unsafeDrop (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) mv (PrimState m) a
v
                Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
              else Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
    )
    Int
1
    Int
nothing
    Int
nothing
  where
    n :: Int
n = mv (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
GM.length mv (PrimState m) a
v
    nothing :: Int
nothing = -Int
1
{-# INLINE nextPermutation #-}

{- |
>>> U.modify (void . prevPermutation) $ U.fromList "acb"
"abc"
>>> U.modify (void . prevPermutation) $ U.fromList "abc"
"abc"
>>> U.modify (void . prevPermutation) $ U.fromList "a"
"a"
>>> U.modify (void . prevPermutation) $ U.fromList ""
""
-}
prevPermutation :: (GM.MVector mv a, Ord a, PrimMonad m) => mv (PrimState m) a -> m Bool
prevPermutation :: forall (mv :: * -> * -> *) a (m :: * -> *).
(MVector mv a, Ord a, PrimMonad m) =>
mv (PrimState m) a -> m Bool
prevPermutation mv (PrimState m) a
v = do
  ((Int -> Int -> Int -> m Bool) -> Int -> Int -> Int -> m Bool)
-> Int -> Int -> Int -> m Bool
forall a. (a -> a) -> a
fix
    ( \Int -> Int -> Int -> m Bool
loop !Int
i !Int
k !Int
l ->
        if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n
          then do
            a
prev <- mv (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead mv (PrimState m) a
v (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
            a
cur <- mv (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead mv (PrimState m) a
v Int
i
            if a
prev a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
cur
              then Int -> Int -> Int -> m Bool
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
i
              else
                if Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
nothing
                  then do
                    a
vk <- mv (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead mv (PrimState m) a
v Int
k
                    if a
vk a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
cur
                      then Int -> Int -> Int -> m Bool
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
k Int
i
                      else Int -> Int -> Int -> m Bool
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
k Int
l
                  else Int -> Int -> Int -> m Bool
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
k Int
l
          else do
            if Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
nothing
              then do
                mv (PrimState m) a -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> Int -> m ()
GM.unsafeSwap mv (PrimState m) a
v Int
k Int
l
                mv (PrimState m) a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> m ()
GM.reverse (mv (PrimState m) a -> m ()) -> mv (PrimState m) a -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> mv (PrimState m) a -> mv (PrimState m) a
forall (v :: * -> * -> *) a s. MVector v a => Int -> v s a -> v s a
GM.unsafeDrop (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) mv (PrimState m) a
v
                Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
              else Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
    )
    Int
1
    Int
nothing
    Int
nothing
  where
    n :: Int
n = mv (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
GM.length mv (PrimState m) a
v
    nothing :: Int
nothing = -Int
1
{-# INLINE prevPermutation #-}

{- |
>>> U.modify (void . flip kthPermutation (120 - 1)) $ U.fromList "abcde"
"edcba"
>>> U.modify (void . flip kthPermutation 0) $ U.fromList "abcde"
"abcde"
>>> U.modify (void . flip kthPermutation 1) $ U.fromList "abcde"
"abced"
>>> U.modify (void . flip kthPermutation 999) $ U.fromList "abcde"
"abcde"
-}
kthPermutation :: (GM.MVector mv a, Ord a, PrimMonad m) => mv (PrimState m) a -> Int -> m Bool
kthPermutation :: forall (mv :: * -> * -> *) a (m :: * -> *).
(MVector mv a, Ord a, PrimMonad m) =>
mv (PrimState m) a -> Int -> m Bool
kthPermutation mv (PrimState m) a
v Int
k0 = do
  if mv (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
GM.length mv (PrimState m) a
v Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
cacheSize Bool -> Bool -> Bool
|| Vector Int
fact Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
U.! mv (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
GM.length mv (PrimState m) a
v Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
k0
    then Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
    else do
      m Int -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m Int -> m ()) -> m Int -> m ()
forall a b. (a -> b) -> a -> b
$
        (Int -> Int -> m Int) -> Int -> Stream m Int -> m Int
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m a) -> a -> Stream m b -> m a
MS.foldlM'
          ( \Int
k Int
i -> do
              let f :: Int
f = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
U.unsafeIndex Vector Int
fact (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                  (Int
q, Int
r) = Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
quotRem Int
k Int
f
              ((Int -> m ()) -> Stream m Int -> m ())
-> Stream m Int -> (Int -> m ()) -> m ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Int -> m ()) -> Stream m Int -> m ()
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Stream m a -> m ()
MS.mapM_ (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Stream m Int
forall (m :: * -> *). Monad m => Int -> Int -> Stream m Int
>.. Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
j -> do
                mv (PrimState m) a -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> Int -> m ()
GM.unsafeSwap mv (PrimState m) a
v Int
j (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
              Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
r
          )
          Int
k0
          (Int
0 Int -> Int -> Stream m Int
forall (m :: * -> *). Monad m => Int -> Int -> Stream m Int
..< Int
n)
      Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
  where
    n :: Int
n = mv (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
GM.length mv (PrimState m) a
v
    cacheSize :: Int
cacheSize = Int
21
    !fact :: Vector Int
fact = (Int -> Int -> Int) -> Int -> Vector Int -> Vector Int
forall a b.
(Unbox a, Unbox b) =>
(a -> b -> a) -> a -> Vector b -> Vector a
U.scanl' Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) Int
1 (Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
U.generate (Int
cacheSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
{-# INLINE kthPermutation #-}