module Test.Agata.Strategies where

import Test.QuickCheck
import Control.Monad.State.Lazy
import Test.Agata.Common


type Strategy a = Int -> Dimension a -> Gen(Improving ())

listStrategy :: (Int -> Dimension a -> Gen [Int]) -> Strategy a
listStrategy f lev0 s = do
    lvls <- f lev0 s
    return $ do
      (lvl,r,[]) <- get
      let d = (lev0 - lvl) + 1
      let k = lvls !! (lvl-1)
      ms <- lift $ piles r k
      put(lvl,0,ms) 


linearSize :: Strategy a
linearSize size _ = return $ do
  (lvl,r,[]) <- get
  k <- lift $ choose (0,size)
  ms <- lift $ piles r k
  put(lvl,0,ms)

quadraticSize :: Strategy a
quadraticSize size lev0 = return $ do
  (lvl,r,[]) <- get
  k <- lift $ choose (0,size*((fromIntegral lev0 - lvl) + 1))
  ms <- lift $ piles r k
  put(lvl,0,ms) 

partitions :: Strategy a
partitions = listStrategy $ \s l -> do
      xs <- sequence $ replicate (fromIntegral l-1) $ choose (0,s)
      permute (s:xs)

exponentialSize :: Strategy a
exponentialSize s _ = return $ do 
  (lvl,r,[]) <- get
  ns <- sequence $ replicate r $ lift $ choose (0,s)
  put (lvl,0,ns)

fixedSize :: Strategy a
fixedSize = listStrategy $ \s l -> piles (fromIntegral l) s
  
randomStrategy :: [Strategy a] -> Strategy a
randomStrategy ls l s = oneof $ map (\f -> f l s) ls