{-# LANGUAGE FlexibleContexts #-}

-- | Internal module for sampling of cycle graph partitions.
-- Import 'RandomCycle.Vector' instead.
module RandomCycle.Vector.Cycle where

import Control.Monad (unless)
import Control.Monad.Primitive (PrimMonad, PrimState, liftPrim)
import Data.STRef
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import System.Random.MWC.Distributions (uniformPermutation, uniformShuffleM)
import System.Random.Stateful

{- INTERNAL -}

-- | Internal. Helper for uniformCyclePartitionThin so as to avoid
-- re-allocating the input vector for each rejected sample.
-- IMPORTANT: Caller's responsibility to ensure proper
-- management of the 'chk' for match found.
uniformCyclePartitionThinM ::
  (StatefulGen g m, PrimMonad m) =>
  STRef (PrimState m) Bool ->
  STRef (PrimState m) Int ->
  ((Int, Int) -> Bool) ->
  V.MVector (PrimState m) Int ->
  g ->
  m ()
uniformCyclePartitionThinM chk maxit r v gen = do
  maxitVal <- liftPrim $ readSTRef maxit

  unless (maxitVal <= 0) $ do
    uniformShuffleM v gen
    -- TODO: Repeated calls to freeze, indexed
    -- a possible opportunity for optimization,
    -- e.g. with imap or a check that takes 'chk'
    -- reference and shortcircuits.
    vVal <- V.freeze v
    if V.all r (V.indexed vVal)
      then do
        liftPrim $ modifySTRef' chk (const True)
      else do
        liftPrim $ modifySTRef' maxit (\x -> x - 1)
        uniformCyclePartitionThinM chk maxit r v gen

-- | Internal. Helper for'uniformCycle'. Caller's responsibility to ensure
-- first input is the length of the vector, and that i ranges over @[0..n-2]@.
-- At the moment, it favors 'MV.swap' over 'MV.unsafeSwap' to maintain bounds
-- checks.
swapIt ::
  (StatefulGen g m, PrimMonad m) =>
  -- | Vector length.
  Int ->
  -- | Vector to modify.
  MV.MVector (PrimState m) Int ->
  -- | Generator.
  g ->
  -- | Index.
  Int ->
  m ()
swapIt n mv g i = do
  let m = n - 2 - i
  j <- uniformRM (0, m) g
  MV.swap mv j (m + 1)

{- RANDOM -}

-- | Implements [Sattolo's algorithm](https://algo.inria.fr/seminars/summary/Wilson2004b.pdf)
-- to sample a full cycle permutation uniformly over /(n-1)!/ possibilities in /O(n)/ time.
-- The algorithm is nearly identical to the Fisher-Yates shuffle on @[0..n-1]@, and therefore
-- this implementation is very similar to that of 'uniformPermutation'.
--
-- This will throw an exception with syntax analogous to 'uniformPermutation'
-- if the provided size is negative.
--
-- ==== __Examples__
--
-- >>> import System.Random.Stateful
-- >>> import RandomCycle.Vector
-- >>> runSTGen_ (mkStdGen 1901) $ uniformCycle 4
-- [(0,3),(1,0),(2,1),(3,2)]
uniformCycle ::
  (StatefulGen g m, PrimMonad m) =>
  -- | Size /n/ of cycle.
  Int ->
  g ->
  m (V.Vector (Int, Int))
uniformCycle n _ | n < 0 = error "RandomCycle.Vector.Cycle: size must be >= 0"
uniformCycle n gen = do
  mv <- MV.generateM n pure
  mapM_ (swapIt n mv gen) [0 .. n - 2]
  V.indexed <$> V.freeze mv

-- | Select a partition of @[0..n-1]@ into disjoint
--  [cycle graphs](https://en.wikipedia.org/wiki/Cycle_graph),
--  uniformly over the /n!/ possibilities. The sampler relies on the fact that such
--  partitions are isomorphic with the permutations of @[0..n-1]@ via the map sending
--  a permutation /\sigma/ to the edge set /\{(i, \sigma(i))\}_0^{n-1}/.
--
--  Therefore, this function simply calls 'uniformPermutation' and tuples the result with its
--  indices. The returned value is a vector of edges. /O(n)/, since 'uniformPermutation'
--  implements the [Fisher-Yates shuffle](https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle).
--
--  'uniformPermutation' uses in-place mutation, so this function must be run in a 'PrimMonad'
--  context.
--
-- ==== __Examples__
--
-- >>> import System.Random.Stateful
-- >>> import qualified RandomCycle.Vector as RV
-- >>> import Data.Vector (Vector)
-- >>> runSTGen_ (mkStdGen 1305) $ RV.uniformCyclePartition 4 :: Vector (Int, Int)
-- [(0,1),(1,3),(2,2),(3,0)]
uniformCyclePartition ::
  (StatefulGen g m, PrimMonad m) =>
  Int ->
  g ->
  m (V.Vector (Int, Int))
uniformCyclePartition n gen = V.indexed <$> uniformPermutation n gen

-- TODO: apply short-circuiting behavior by creating  modification
-- of 'uniformSuffleM' that carries a validity state and short-circuits as
-- soon as some edge does not satisfy the predicate.
-- current implementation is the lazy one (as in human-lazy). note that would require
-- posting a notice in this module in accordance with the BSD2 license of mwc-random.

-- | Uniform selection of a cycle partition graph of @[0..n-1]@ elements,
-- conditional on an edge-wise predicate. See 'uniformCyclePartition' for
-- details on the sampler.
--
-- /O(n\/p)/, where /p/ is the probability that a uniformly sampled
-- cycle partition graph (over all /n!/ possible) satisfies the conditions.
-- This can be highly non-linear since /p/ in general is a function of /n/.
--
-- Since this is a rejection sampling method, the user is asked to provide
-- a counter for the maximum number of sampling attempts in order to guarantee
-- termination in cases where the edge predicate has probability of success close
-- to zero.
--
-- Note this will return @pure Nothing@ if given a number of vertices that is
-- non-positive, in the third argument, unlike 'uniformCyclePartition' which
-- will throw an error.
--
-- ==== __Examples__
--
-- >>> import System.Random.Stateful
-- >>> import qualified RandomCycle.Vector as RV
-- >>> import Data.Vector (Vector)
-- >>> -- No self-loops
-- >>> rule = uncurry (/=)
-- >>> n = 5
-- >>> maxit = n * 1000
-- >>> runSTGen_ (mkStdGen 3) $ RV.uniformCyclePartitionThin maxit rule n
-- Just [(0,2),(1,3),(2,0),(3,4),(4,1)]
uniformCyclePartitionThin ::
  (StatefulGen g m, PrimMonad m) =>
  -- | maximum number of draws to attempt
  Int ->
  -- | edge-wise predicate, which all edges in the result must satisfy
  ((Int, Int) -> Bool) ->
  -- | number of vertices, which will be labeled @[0..n-1]@
  Int ->
  g ->
  m (Maybe (V.Vector (Int, Int)))
uniformCyclePartitionThin maxit _ n _en | maxit <= 0 || n <= 0 = pure Nothing
uniformCyclePartitionThin maxit r n gen = do
  let v = V.generate n id
  mv <- V.thaw v
  chk' <- liftPrim $ newSTRef False
  maxit' <- liftPrim $ newSTRef maxit

  uniformCyclePartitionThinM chk' maxit' r mv gen

  chk <- liftPrim $ readSTRef chk'
  if chk
    then do
      Just . V.indexed <$> V.freeze mv
    else pure Nothing
