{-# LANGUAGE LambdaCase #-}

-- |
-- Most of the code is borrowed from
-- <http://haskell.1045720.n5.nabble.com/darcs-patch-GenT-monad-transformer-variant-of-Gen-QuickCheck-2-td3172136.html a mailing list discussion>.
-- Therefor, credits go to Paul Johnson and Felix Martini.
module QuickCheck.GenT
  ( GenT,
    runGenT,
    MonadGen (..),

    -- * Lifted functions
    arbitrary',
    oneof,
    frequency,
    elements,
    growingElements,
    getSize,
    scale,
    suchThat,
    suchThatMap,
    suchThatMaybe,
    applyArbitrary2,
    applyArbitrary3,
    applyArbitrary4,
    listOf,
    listOf1,
    vectorOf,
    vector,
    infiniteListOf,
    infiniteList,
    shuffle,
    sublistOf,
    orderedList,

    -- * Re-exports
    Arbitrary (..),
    QC.Gen,

    -- * Safe functions
    oneofMay,
    elementsMay,
    growingElementsMay,
  )
where

import QuickCheck.GenT.Prelude
import qualified System.Random as Random
import Test.QuickCheck (Arbitrary (..))
import qualified Test.QuickCheck.Arbitrary as QC
import qualified Test.QuickCheck.Gen as QC
import qualified Test.QuickCheck.Random as QC

newtype GenT m a = GenT {unGenT :: QC.QCGen -> Int -> m a}

instance MFunctor GenT where
  hoist f (GenT g) = GenT $ \r n -> f $ g r n

instance (Functor m) => Functor (GenT m) where
  fmap f m = GenT $ \r n -> fmap f $ unGenT m r n

instance (Monad m) => Monad (GenT m) where
  return = pure
  m >>= k = GenT $ \r n -> do
    let (r1, r2) = Random.split r
    a <- unGenT m r1 n
    unGenT (k a) r2 n

instance (MonadFail m) => MonadFail (GenT m) where
  fail msg = GenT (\_ _ -> fail msg)

instance (Functor m, Monad m) => Applicative (GenT m) where
  pure a = GenT (\_ _ -> return a)
  (<*>) = ap

instance MonadTrans GenT where
  lift m = GenT (\_ _ -> m)

instance (MonadIO m) => MonadIO (GenT m) where
  liftIO = lift . liftIO

class (Applicative g, Monad g) => MonadGen g where
  liftGen :: QC.Gen a -> g a
  variant :: (Integral n) => n -> g a -> g a
  sized :: (Int -> g a) -> g a
  resize :: Int -> g a -> g a
  choose :: (Random.Random a) => (a, a) -> g a

instance (Applicative m, Monad m) => MonadGen (GenT m) where
  liftGen gen = GenT $ \r n -> return $ QC.unGen gen r n
  choose rng = GenT $ \r _ -> return $ fst $ Random.randomR rng r
  variant k (GenT g) = GenT $ \r n -> g (var k r) n
  sized f = GenT $ \r n -> let GenT g = f n in g r n
  resize n (GenT g) = GenT $ \r _ -> g r n

instance MonadGen QC.Gen where
  liftGen = id
  variant k (QC.MkGen g) = QC.MkGen $ \r n -> g (var k r) n
  sized f = QC.MkGen $ \r n -> let QC.MkGen g = f n in g r n
  resize n (QC.MkGen g) = QC.MkGen $ \r _ -> g r n
  choose range = QC.MkGen $ \r _ -> fst $ Random.randomR range r

runGenT :: GenT m a -> QC.Gen (m a)
runGenT (GenT run) = QC.MkGen run

-- |
-- Private variant-generating function.  Converts an integer into a chain
-- of (fst . split) and (snd . split) applications.  Every integer (including
-- negative ones) will give rise to a different random number generator in
-- log2 n steps.
var :: (Integral n) => n -> QC.QCGen -> QC.QCGen
var k =
  (if k == k' then id else var k') . (if even k then fst else snd) . Random.split
  where
    k' = k `div` 2

-- ** Lifted functions

arbitrary' :: (Arbitrary a, MonadGen m) => m a
arbitrary' = liftGen arbitrary

getSize :: (MonadGen m) => m Int
getSize = liftGen getSize

scale :: (MonadGen m) => (Int -> Int) -> m a -> m a
scale f g = sized (\n -> resize (f n) g)

applyArbitrary2 :: (MonadGen m) => (Arbitrary a, Arbitrary b) => (a -> b -> r) -> m r
applyArbitrary2 = liftGen . QC.applyArbitrary2

applyArbitrary3 :: (MonadGen m) => (Arbitrary a, Arbitrary b, Arbitrary c) => (a -> b -> c -> r) -> m r
applyArbitrary3 = liftGen . QC.applyArbitrary3

applyArbitrary4 :: (MonadGen m) => (Arbitrary a, Arbitrary b, Arbitrary c, Arbitrary d) => (a -> b -> c -> d -> r) -> m r
applyArbitrary4 = liftGen . QC.applyArbitrary4

infiniteListOf :: (MonadGen m) => m a -> m [a]
infiniteListOf = sequence . repeat

infiniteList :: (Arbitrary a, MonadGen m) => m [a]
infiniteList = infiniteListOf arbitrary'

shuffle :: (MonadGen m) => [a] -> m [a]
shuffle = liftGen . QC.shuffle

sublistOf :: (MonadGen m) => [a] -> m [a]
sublistOf = liftGen . QC.sublistOf

orderedList :: (Ord a, Arbitrary a, MonadGen m) => m [a]
orderedList = liftGen QC.orderedList

-- ** Common generator combinators

-- | Generates a value that satisfies a predicate.
suchThat :: (MonadGen m) => m a -> (a -> Bool) -> m a
gen `suchThat` p =
  do
    mx <- gen `suchThatMaybe` p
    case mx of
      Just x -> return x
      Nothing -> sized (\n -> resize (n + 1) (gen `suchThat` p))

-- | Generates a value for which the given function returns a 'Just', and then
-- applies the function.
suchThatMap :: (MonadGen m) => m a -> (a -> Maybe b) -> m b
gen `suchThatMap` f =
  fmap fromJust $ fmap f gen `suchThat` isJust

-- | Tries to generate a value that satisfies a predicate.
suchThatMaybe :: (MonadGen m) => m a -> (a -> Bool) -> m (Maybe a)
gen `suchThatMaybe` p = sized (try 0 . max 1)
  where
    try _ 0 = return Nothing
    try k n = do
      x <- resize (2 * k + n) gen
      if p x then return (Just x) else try (k + 1) (n - 1)

-- | Generates a list of random length. The maximum length depends on the
-- size parameter.
listOf :: (MonadGen m) => m a -> m [a]
listOf gen = sized $ \n ->
  do
    k <- choose (0, n)
    vectorOf k gen

-- | Generates a non-empty list of random length. The maximum length
-- depends on the size parameter.
listOf1 :: (MonadGen m) => m a -> m [a]
listOf1 gen = sized $ \n ->
  do
    k <- choose (1, 1 `max` n)
    vectorOf k gen

-- | Generates a list of the given length.
vectorOf :: (MonadGen m) => Int -> m a -> m [a]
vectorOf k gen = sequence [gen | _ <- [1 .. k]]

-- | Generates a list of a given length.
vector :: (Arbitrary a, MonadGen m) => Int -> m [a]
vector n = vectorOf n arbitrary'

-- * Partial functions

-- | Randomly uses one of the given generators. The input list
-- must be non-empty.
oneof :: (MonadGen m) => [m a] -> m a
oneof =
  fmap (fromMaybe (error "QuickCheck.GenT.oneof used with empty list"))
    . oneofMay

-- | Chooses one of the given generators, with a weighted random distribution.
-- The input list must be non-empty.
frequency :: (MonadGen m) => [(Int, m a)] -> m a
frequency [] = error "QuickCheck.GenT.frequency used with empty list"
frequency xs0 = choose (1, tot) >>= (`pick` xs0)
  where
    tot = sum (map fst xs0)

    pick n ((k, x) : xs)
      | n <= k = x
      | otherwise = pick (n - k) xs
    pick _ _ = error "QuickCheck.GenT.pick used with empty list"

-- | Generates one of the given values. The input list must be non-empty.
elements :: (MonadGen m) => [a] -> m a
elements =
  fmap (fromMaybe (error "QuickCheck.GenT.elements used with empty list"))
    . elementsMay

-- | Takes a list of elements of increasing size, and chooses
-- among an initial segment of the list. The size of this initial
-- segment increases with the size parameter.
-- The input list must be non-empty.
growingElements :: (MonadGen m) => [a] -> m a
growingElements =
  fmap (fromMaybe (error "QuickCheck.GenT.growingElements used with empty list"))
    . growingElementsMay

-- * Total functions resulting in Maybe

-- |
-- Randomly uses one of the given generators.
oneofMay :: (MonadGen m) => [m a] -> m (Maybe a)
oneofMay = \case
  [] -> return Nothing
  l -> fmap Just $ choose (0, length l - 1) >>= (l !!)

-- | Generates one of the given values.
elementsMay :: (MonadGen m) => [a] -> m (Maybe a)
elementsMay = \case
  [] -> return Nothing
  l -> Just . (l !!) <$> choose (0, length l - 1)

-- | Takes a list of elements of increasing size, and chooses
-- among an initial segment of the list. The size of this initial
-- segment increases with the size parameter.
growingElementsMay :: (MonadGen m) => [a] -> m (Maybe a)
growingElementsMay = \case
  [] -> return Nothing
  xs -> fmap Just $ sized $ \n -> elements (take (1 `max` size n) xs)
    where
      k = length xs
      mx = 100
      log' = round . log . fromIntegral
      size n = (log' n + 1) * k `div` log' mx