{-# LANGUAGE NoMonomorphismRestriction, RecordWildCards
  , StandaloneDeriving, MultiParamTypeClasses, FunctionalDependencies
  , TypeSynonymInstances, FlexibleInstances, FlexibleContexts
  , UndecidableInstances, DeriveDataTypeable, GeneralizedNewtypeDeriving
  , OverlappingInstances, RankNTypes, PatternGuards
  #-}

module Codec.TPTP.QuickCheck where
    
import Test.QuickCheck
import Control.Monad
import Control.Applicative
import Data.Char
import Data.Array.ST
import Data.Array.IArray
    
argsFreq :: (Int -> Gen a) -> Gen a
argsFreq f = frequency [ (10,f 0)
                        , (10,f 1)
                        , (10,f 2)
                        , (2 ,f 3)
                        , (2 ,f 4)
                        , (1, f 15)
                        ]
              
arbVar :: Gen [Char]
arbVar = liftM2 (:) (elements "WXZY") (show <$> choose(1::Int,3))  
        
-- sorry I don't feel like calculating this distribution deductively at the moment ;)
-- | Return random partition of n items into the given number of buckets if buckets>0
--
-- Return [] if buckets=0
arbPartition :: Int -> Int -> Gen [Int]
arbPartition 0 _ = return []
arbPartition 1 n = return [n]
arbPartition buckets n = do
  choices <- replicateM n (choose (1,buckets))

  return . Data.Array.IArray.elems . runSTUArray $
   do
    arrr <- newArray (1,buckets) 0
    forM_ choices (\bucket -> writeArray arrr bucket . succ =<< readArray arrr bucket)
    return arrr
           
arbPrintable :: Gen [Char]
arbPrintable = listOf (arbitrary `suchThat` printable)
               
printable :: Char -> Bool
printable x = isAscii x && isPrint x --isDigit x || isLetter x || x `elem` " _+!@#$%^&*()[]{}-=?.,;'\":/\\"
    
arbLowerWord :: Gen String
arbLowerWord = (:) `fmap` elements ['a'..'z'] `ap` listOf (elements (['a'..'z']++['A'..'Z']++"_"))


arbUpperWord :: Gen String
arbUpperWord = (:) `fmap` elements ['A'..'Z'] `ap` listOf (elements (['a'..'z']++['A'..'Z']++"_"))

               
arbNum :: forall a a1. (Arbitrary a, Num a) => (a -> a1) -> Gen a1
arbNum f =
    frequency [(1,fmap f arbitrary)
              ,(8, fmap (f . fromIntegral) (arbitrary::Gen Int))
              ]