{-# LANGUAGE FlexibleContexts #-}

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

import Control.Monad (when)
import Control.Monad.Primitive (PrimMonad, PrimState, liftPrim)
import Data.STRef
import qualified Data.Vector as V
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

  when (maxitVal <= 0) (pure ())

  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
  pure ()

{- RANDOM -}

-- TODO: uniform (full) cycle with [Sattolo's algorithm](https://algo.inria.fr/seminars/summary/Wilson2004b.pdf)
-- uniformCycle

-- | 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}\). In other words,
--  the cycle partition graphs are isomorphic with the rotation matrices.
--
--  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 RandomCycle.Vector
-- >>> 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.
uniformCyclePartitionThin ::
  (StatefulGen g m, PrimMonad m) =>
  Int ->
  ((Int, Int) -> Bool) ->
  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