{-# 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 Data.Set qualified 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 i -> i -> Bool
forall a. Ord a => a -> a -> Bool
<= i
b
    then NonEmptyRange i -> Range i
forall i. NonEmptyRange i -> Range i
RangeNonEmpty (i -> Natural -> NonEmptyRange i
forall i. i -> Natural -> NonEmptyRange i
NonEmptyRange i
a (i -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral (i
b i -> i -> i
forall a. Num a => a -> a -> a
- i
a)))
    else Range i
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 -> Set i -> m (Set i)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Set i
forall a. Set a
Set.empty
  RangeNonEmpty NonEmptyRange i
r ->
    (RangeWithGaps i -> Set i) -> m (RangeWithGaps i) -> m (Set i)
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap RangeWithGaps i -> Set i
forall i. RangeWithGaps i -> Set i
gaps (m (RangeWithGaps i) -> m (Set i))
-> m (RangeWithGaps i) -> m (Set i)
forall a b. (a -> b) -> a -> b
$
      (StateT (RangeWithGaps i) m ()
 -> RangeWithGaps i -> m (RangeWithGaps i))
-> RangeWithGaps i
-> StateT (RangeWithGaps i) m ()
-> m (RangeWithGaps i)
forall a b c. (a -> b -> c) -> b -> a -> c
flip StateT (RangeWithGaps i) m ()
-> RangeWithGaps i -> m (RangeWithGaps i)
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m s
execStateT (NonEmptyRange i -> Set i -> RangeWithGaps i
forall i. NonEmptyRange i -> Set i -> RangeWithGaps i
RangeWithGaps NonEmptyRange i
r Set i
forall a. Set a
Set.empty) (StateT (RangeWithGaps i) m () -> m (RangeWithGaps i))
-> StateT (RangeWithGaps i) m () -> m (RangeWithGaps i)
forall a b. (a -> b) -> a -> b
$
        [Natural]
-> (Natural -> StateT (RangeWithGaps i) m ())
-> StateT (RangeWithGaps i) m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Natural
1 .. Natural
n] ((Natural -> StateT (RangeWithGaps i) m ())
 -> StateT (RangeWithGaps i) m ())
-> (Natural -> StateT (RangeWithGaps i) m ())
-> StateT (RangeWithGaps i) m ()
forall a b. (a -> b) -> a -> b
$ \Natural
_ -> do
          StateT (RangeWithGaps i) m (RangeWithGaps i)
forall s (m :: * -> *). MonadState s m => m s
get StateT (RangeWithGaps i) m (RangeWithGaps i)
-> (RangeWithGaps i
    -> StateT (RangeWithGaps i) m (RangeWithGaps i))
-> StateT (RangeWithGaps i) m (RangeWithGaps i)
forall a b.
StateT (RangeWithGaps i) m a
-> (a -> StateT (RangeWithGaps i) m b)
-> StateT (RangeWithGaps i) m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (m (RangeWithGaps i) -> StateT (RangeWithGaps i) m (RangeWithGaps i)
forall (m :: * -> *) a.
Monad m =>
m a -> StateT (RangeWithGaps i) m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m (RangeWithGaps i)
 -> StateT (RangeWithGaps i) m (RangeWithGaps i))
-> (RangeWithGaps i -> m (RangeWithGaps i))
-> RangeWithGaps i
-> StateT (RangeWithGaps i) m (RangeWithGaps i)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RangeWithGaps i -> m (RangeWithGaps i)
forall (m :: * -> *) i.
(MonadRandom m, Random i, Integral i) =>
RangeWithGaps i -> m (RangeWithGaps i)
randomlyRemove) StateT (RangeWithGaps i) m (RangeWithGaps i)
-> (RangeWithGaps i -> StateT (RangeWithGaps i) m ())
-> StateT (RangeWithGaps i) m ()
forall a b.
StateT (RangeWithGaps i) m a
-> (a -> StateT (RangeWithGaps i) m b)
-> StateT (RangeWithGaps i) m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= RangeWithGaps i -> StateT (RangeWithGaps i) m ()
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 =
  RangeWithGaps i -> m (Maybe i)
forall (m :: * -> *) i.
(MonadRandom m, Random i, Integral i) =>
RangeWithGaps i -> m (Maybe i)
randomFromRangeWithGaps RangeWithGaps i
rg m (Maybe i) -> (Maybe i -> RangeWithGaps i) -> m (RangeWithGaps i)
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.insert i (gaps 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 :: forall i. RangeWithGaps i -> NonEmptyRange i
rangeWithoutGaps :: NonEmptyRange i
rangeWithoutGaps, Set i
gaps :: forall i. RangeWithGaps i -> Set i
gaps :: Set i
gaps} = RangeWithGaps i
rg
    NonEmptyRange {i
inclusiveMinBound :: forall i. NonEmptyRange i -> i
inclusiveMinBound :: i
inclusiveMinBound, Natural
offset :: forall i. NonEmptyRange i -> Natural
offset :: Natural
offset} = NonEmptyRange i
rangeWithoutGaps
  in
    if Int -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Set i -> Int
forall a. Set a -> Int
Set.size Set i
gaps) Natural -> Natural -> Bool
forall a. Eq a => a -> a -> Bool
== Natural
offset Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ Natural
1
      then Maybe i -> m (Maybe i)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe i
forall a. Maybe a
Nothing
      else
        i -> Maybe i
forall a. a -> Maybe a
Just
          (i -> Maybe i) -> m i -> m (Maybe i)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> do
            i
r <-
              (i
inclusiveMinBound +)
                (i -> i) -> m i -> m i
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (i, i) -> m i
forall a. Random a => (a, a) -> m a
forall (m :: * -> *) a. (MonadRandom m, Random a) => (a, a) -> m a
getRandomR (i
0, Natural -> i
forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
offset i -> i -> i
forall a. Num a => a -> a -> a
- Int -> i
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Set i -> Int
forall a. Set a -> Int
Set.size Set i
gaps))
            i -> m i
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (i -> m i) -> i -> m i
forall a b. (a -> b) -> a -> b
$ (forall s. ST s i) -> i
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s i) -> i) -> (forall s. ST s i) -> i
forall a b. (a -> b) -> a -> b
$ do
              STRef s i
xRef <- i -> ST s (STRef s i)
forall a s. a -> ST s (STRef s a)
newSTRef i
r
              STRef s [i]
gapQueue <- [i] -> ST s (STRef s [i])
forall a s. a -> ST s (STRef s a)
newSTRef ([i] -> ST s (STRef s [i])) -> [i] -> ST s (STRef s [i])
forall a b. (a -> b) -> a -> b
$ Set i -> [i]
forall a. Set a -> [a]
Set.toAscList Set i
gaps
              let go :: ST s i
go = do
                    i
x <- STRef s i -> ST s i
forall s a. STRef s a -> ST s a
readSTRef STRef s i
xRef
                    STRef s [i] -> ST s [i]
forall s a. STRef s a -> ST s a
readSTRef STRef s [i]
gapQueue ST s [i] -> ([i] -> ST s i) -> ST s i
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
                      i
g : [i]
gs | i
g i -> i -> Bool
forall a. Ord a => a -> a -> Bool
<= i
x -> do
                        STRef s i -> i -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s i
xRef (i
x i -> i -> i
forall a. Num a => a -> a -> a
+ i
1)
                        STRef s [i] -> [i] -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s [i]
gapQueue [i]
gs
                        ST s i
go
                      [i]
_ -> i -> ST s i
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure i
x
              ST s i
go