{-# LANGUAGE BangPatterns #-}

module Algorithms.Random.Shuffle.Pure
  ( shuffle
  , sampleOne
  , sampleSplitOne
  , GetRandR
  ) where


import qualified Data.MonoTraversable  as MT
import qualified Data.Sequences        as S
import qualified System.Random.Shuffle as RS


-- | Monadic action generating an index number for shuffling.
--   The type parameter 'm' is usually some @Monad@.
--   The first argument is a inclusive range /(lo, hi)/.
--   So the retuned value should be /lo <= x <= hi/
--   (like 'System.Random.randomR' and 'System.Random.MWC.uniformR').
type GetRandR m = (Int, Int) -> m Int


-- | Reimplementation of 'RS.shuffleM' in terms of a raw @Monad m => m a@.
--   TODO: Generalize by mono-traversable
shuffle :: Monad m => GetRandR m -> [a] -> m [a]
shuffle _ [] = return []
shuffle randR xs =
  RS.shuffle xs <$> rseqM (length xs - 1)
 where
  rseqM 0 = return []
  rseqM i = (:) <$> randR (0, i) <*> rseqM (i - 1)


-- | Implementation of <https://en.wikipedia.org/wiki/Reservoir_sampling Reservoir sampling>
--   for a single sample.
sampleOne :: (S.IsSequence seq, Monad m) => GetRandR m -> seq -> m (Maybe (MT.Element seq))
sampleOne randR xs =
  case S.uncons xs of
    Just (first, left) -> Just <$> sampleOneDef randR first left
    _                  -> return Nothing


sampleOneDef :: (S.IsSequence seq, Monad m) => GetRandR m -> MT.Element seq -> seq -> m (MT.Element seq)
sampleOneDef randR def = fmap snd . MT.ofoldlM f (size, def)
 where
  f (!i, !current) another = do
    j <- randR (0, i)
    let new = if j == 0 then another else current
    return (i + 1, new)
  size = 1


-- | Implementation of <https://en.wikipedia.org/wiki/Reservoir_sampling Reservoir sampling>
--   for a single sample. By contrast to 'sampleOne',
--   this function returns not chosen elements in addition.
sampleSplitOne :: (S.IsSequence seq, Monad m) => GetRandR m -> seq -> m (Maybe (MT.Element seq, seq))
sampleSplitOne randR xs =
  case S.uncons xs of
    Just (first, left) -> Just <$> sampleSplitOneDef randR first left
    _                  -> return Nothing


sampleSplitOneDef :: (S.IsSequence seq, Monad m) => GetRandR m -> MT.Element seq -> seq -> m (MT.Element seq, seq)
sampleSplitOneDef randR def = fmap snd . MT.ofoldlM f (size, (def, mempty))
 where
  f (!i, (!current, !notChosens)) another = do
    j <- randR (0, i)
    let new =
          if j == 0
            then (another, current `S.cons` notChosens)
            else (current, another `S.cons` notChosens)
    return (i + 1, new)
  size = 1