{-# language Rank2Types #-}
{-# language ScopedTypeVariables #-}
module EasyTest.Generators
  ( -- * Generators
    random
  , random'
  , bool
  , word8
  , char
  , int
  , double
  , word
  , int'
  , char'
  , double'
  , word'
  , word8'
  , pick
  , listOf
  , listsOf
  , pair
  , mapOf
  , mapsOf
  ) where

import Control.Applicative
import Control.Concurrent.STM
import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.Reader
import Data.Map (Map)
import Data.Word
import System.Random (Random)
import qualified Data.Map as Map
import qualified System.Random as Random

import EasyTest.Internal

-- | Generate a random value
random :: forall a. Random a => Test a
random = do
  rng <- asks envRng
  liftIO . atomically $ do
    rng0 <- readTVar rng
    let (a :: a, rng1) = Random.random rng0
    writeTVar rng rng1
    pure a

-- | Generate a bounded random value. Inclusive on both sides.
random' :: Random a => a -> a -> Test a
random' lower upper = do
  rng <- asks envRng
  liftIO . atomically $ do
    rng0 <- readTVar rng
    let (a, rng1) = Random.randomR (lower,upper) rng0
    writeTVar rng rng1
    pure a

bool :: Test Bool
bool = random

word8 :: Test Word8
word8 = random

-- | Generate a random 'Char'
char :: Test Char
char = random

-- | Generate a random 'Int'
int :: Test Int
int = random

-- | Generate a random 'Double'
double :: Test Double
double = random

-- | Generate a random 'Word'
word :: Test Word
word = random

-- | Generate a random 'Int' in the given range
-- Note: @int' 0 5@ includes both @0@ and @5@
int' :: Int -> Int -> Test Int
int' = random'

-- | Generate a random 'Char' in the given range
-- Note: @char' 'a' 'z'@ includes both @'a'@ and @'z'@.
char' :: Char -> Char -> Test Char
char' = random'

-- | Generate a random 'Double' in the given range
-- Note: @double' 0 1@ includes both @0@ and @1@.
double' :: Double -> Double -> Test Double
double' = random'

-- | Generate a random 'Double' in the given range
-- Note: @word' 0 10@ includes both @0@ and @10@.
word' :: Word -> Word -> Test Word
word' = random'

-- | Generate a random 'Double' in the given range
-- Note: @word8' 0 10@ includes both @0@ and @10@.
word8' :: Word8 -> Word8 -> Test Word8
word8' = random'

-- | Sample uniformly from the given list of possibilities
pick :: [a] -> Test a
pick as = let n = length as; ind = picker n as in do
  i <- int' 0 (n - 1)
  Just a <- pure (ind i)
  pure a

picker :: Int -> [a] -> (Int -> Maybe a)
picker _ [] = const Nothing
picker _ [a] = \i -> if i == 0 then Just a else Nothing
picker size as = go where
  lsize = size `div` 2
  rsize = size - lsize
  (l,r) = splitAt lsize as
  lpicker = picker lsize l
  rpicker = picker rsize r
  go i = if i < lsize then lpicker i else rpicker (i - lsize)

-- | Alias for 'replicateM'
listOf :: Int -> Test a -> Test [a]
listOf = replicateM

-- | Generate a list of lists of the given sizes,
-- an alias for @sizes \`forM\` \\n -> listOf n gen@
listsOf :: [Int] -> Test a -> Test [[a]]
listsOf sizes gen = sizes `forM` \n -> listOf n gen

-- | Alias for @liftA2 (,)@.
pair :: Test a -> Test b -> Test (a,b)
pair = liftA2 (,)

-- | Generate a @Data.Map k v@ of the given size.
mapOf :: Ord k => Int -> Test k -> Test v -> Test (Map k v)
mapOf n k v = Map.fromList <$> listOf n (pair k v)

-- | Generate a @[Data.Map k v]@ of the given sizes.
mapsOf :: Ord k => [Int] -> Test k -> Test v -> Test [Map k v]
mapsOf sizes k v = sizes `forM` \n -> mapOf n k v