{-# LANGUAGE BangPatterns #-}

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


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

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