{-# LANGUAGE NamedFieldPuns #-}

module Freckle.App.Random
  ( smallRandomSubsetOfLargeIntegerRange
  , Range (..)
  , NonEmptyRange (..)
  , inclusiveRange
  ) where

import Freckle.App.Prelude

import Control.Monad.Random (MonadRandom (..), Random)
import Control.Monad.ST (runST)
import Control.Monad.State.Strict (execStateT, get, put)
import Data.Functor ((<&>))
import Data.STRef (newSTRef, readSTRef, writeSTRef)
import Numeric.Natural (Natural)

import qualified Data.Set as Set

-- | A possibly-empty contiguous range of integers
data Range i
  = RangeEmpty
  | RangeNonEmpty (NonEmptyRange i)

-- | A nonempty contiguous range of integers
data NonEmptyRange i = NonEmptyRange
  { forall i. NonEmptyRange i -> i
inclusiveMinBound :: i
  , forall i. NonEmptyRange i -> Natural
offset :: Natural
  -- ^ The size of the range minus one
  }

inclusiveRange
  :: Integral i
  => i
  -- ^ Lower bound, inclusive
  -> i
  -- ^ Upper bound, inclusive
  -> Range i
inclusiveRange :: forall i. Integral i => i -> i -> Range i
inclusiveRange i
a i
b =
  if i
a forall a. Ord a => a -> a -> Bool
<= i
b
    then forall i. NonEmptyRange i -> Range i
RangeNonEmpty (forall i. i -> Natural -> NonEmptyRange i
NonEmptyRange i
a (forall a b. (Integral a, Num b) => a -> b
fromIntegral (i
b forall a. Num a => a -> a -> a
- i
a)))
    else forall i. Range i
RangeEmpty

-- | Select a fixed number of items uniformly at random
--   from a contiguous range of integers
--
-- This process accommodates selecting from a large range, but only has
-- reasonable performance when the number of items being selected is small
-- (it is quadratic in the number of items).
--
-- If the requested size is greater than or equal to the range, the entire
-- range is returned.
--
-- e.g. @smallRandomSubsetOfLargeIntegerRange 10 (inclusiveRange 30 70)@
-- may produce something like @fromList [32,34,45,54,56,58,62,63,64,65]@.
smallRandomSubsetOfLargeIntegerRange
  :: (MonadRandom m, Random i, Integral i)
  => Natural
  -- ^ How many items are wanted
  -> Range i
  -> m (Set i)
smallRandomSubsetOfLargeIntegerRange :: forall (m :: * -> *) i.
(MonadRandom m, Random i, Integral i) =>
Natural -> Range i -> m (Set i)
smallRandomSubsetOfLargeIntegerRange Natural
n = \case
  Range i
RangeEmpty -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Set a
Set.empty
  RangeNonEmpty NonEmptyRange i
r ->
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall i. RangeWithGaps i -> Set i
gaps forall a b. (a -> b) -> a -> b
$
      forall a b c. (a -> b -> c) -> b -> a -> c
flip forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m s
execStateT (forall i. NonEmptyRange i -> Set i -> RangeWithGaps i
RangeWithGaps NonEmptyRange i
r forall a. Set a
Set.empty) forall a b. (a -> b) -> a -> b
$
        forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Natural
1 .. Natural
n] forall a b. (a -> b) -> a -> b
$ \Natural
_ -> do
          forall s (m :: * -> *). MonadState s m => m s
get forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) i.
(MonadRandom m, Random i, Integral i) =>
RangeWithGaps i -> m (RangeWithGaps i)
randomlyRemove) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall s (m :: * -> *). MonadState s m => s -> m ()
put

data RangeWithGaps i = RangeWithGaps
  { forall i. RangeWithGaps i -> NonEmptyRange i
rangeWithoutGaps :: NonEmptyRange i
  , forall i. RangeWithGaps i -> Set i
gaps :: Set i
  -- ^ The set of items that has been removed from the larger range
  }

-- | Randomly remove an item from a 'RangeWithGaps'.
--
-- This selects uniformly at random an item from a 'RangeWithGaps' and
-- removes it (adds it to the 'gaps' set).
--
-- If every item in the range has already been removed, this does nothing.
randomlyRemove
  :: (MonadRandom m, Random i, Integral i) => RangeWithGaps i -> m (RangeWithGaps i)
randomlyRemove :: forall (m :: * -> *) i.
(MonadRandom m, Random i, Integral i) =>
RangeWithGaps i -> m (RangeWithGaps i)
randomlyRemove RangeWithGaps i
rg =
  forall (m :: * -> *) i.
(MonadRandom m, Random i, Integral i) =>
RangeWithGaps i -> m (Maybe i)
randomFromRangeWithGaps RangeWithGaps i
rg forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
    Maybe i
Nothing -> RangeWithGaps i
rg
    Just i
i -> RangeWithGaps i
rg {gaps :: Set i
gaps = forall a. Ord a => a -> Set a -> Set a
Set.insert i
i (forall i. RangeWithGaps i -> Set i
gaps RangeWithGaps i
rg)}

-- | Randomly select an item from a 'RangeWithGaps'
--
-- This selects uniformly at random an item from the range that is not
-- present in the 'gaps' set.
randomFromRangeWithGaps
  :: (MonadRandom m, Random i, Integral i)
  => RangeWithGaps i
  -> m (Maybe i)
randomFromRangeWithGaps :: forall (m :: * -> *) i.
(MonadRandom m, Random i, Integral i) =>
RangeWithGaps i -> m (Maybe i)
randomFromRangeWithGaps RangeWithGaps i
rg =
  let
    RangeWithGaps {NonEmptyRange i
rangeWithoutGaps :: NonEmptyRange i
rangeWithoutGaps :: forall i. RangeWithGaps i -> NonEmptyRange i
rangeWithoutGaps, Set i
gaps :: Set i
gaps :: forall i. RangeWithGaps i -> Set i
gaps} = RangeWithGaps i
rg
    NonEmptyRange {i
inclusiveMinBound :: i
inclusiveMinBound :: forall i. NonEmptyRange i -> i
inclusiveMinBound, Natural
offset :: Natural
offset :: forall i. NonEmptyRange i -> Natural
offset} = NonEmptyRange i
rangeWithoutGaps
  in
    if forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Set a -> Int
Set.size Set i
gaps) forall a. Eq a => a -> a -> Bool
== Natural
offset forall a. Num a => a -> a -> a
+ Natural
1
      then forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
      else
        forall a. a -> Maybe a
Just
          forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> do
            i
r <-
              (i
inclusiveMinBound forall a. Num a => a -> a -> a
+)
                forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. (MonadRandom m, Random a) => (a, a) -> m a
getRandomR (i
0, forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
offset forall a. Num a => a -> a -> a
- forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Set a -> Int
Set.size Set i
gaps))
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
              STRef s i
xRef <- forall a s. a -> ST s (STRef s a)
newSTRef i
r
              STRef s [i]
gapQueue <- forall a s. a -> ST s (STRef s a)
newSTRef forall a b. (a -> b) -> a -> b
$ forall a. Set a -> [a]
Set.toAscList Set i
gaps
              let go :: ST s i
go = do
                    i
x <- forall s a. STRef s a -> ST s a
readSTRef STRef s i
xRef
                    forall s a. STRef s a -> ST s a
readSTRef STRef s [i]
gapQueue forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
                      i
g : [i]
gs | i
g forall a. Ord a => a -> a -> Bool
<= i
x -> do
                        forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s i
xRef (i
x forall a. Num a => a -> a -> a
+ i
1)
                        forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s [i]
gapQueue [i]
gs
                        ST s i
go
                      [i]
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure i
x
              ST s i
go