-- 
-- (c) Susumu Katayama
--
{-
rewrite of QuickCheck.Arbitrary in the form specialized for each type
@inproceedings{QuickCheck,
        AUTHOR = "Koen Claessen and John Hughes",
        TITLE  = "{QuickCheck}: a lightweight tool for random testing of {Haskell} programs",
        BOOKTITLE = "ICFP'00: Proceedings of the 5th ACM SIGPLAN International Conference on Functional Programming",
        PAGES  = "268-279",
        ORGANIZATION = "ACM",
        YEAR = 2000 }
The original source is released under BSD-style license.
I (Susumu) reimplemented this because QuickCheck-1 had (and has?) some bugs and QuickCheck-2 was not released, but 
maybe I could import and reuse definitions of Arbitrary of QuickCheck-2.
(But still I am interested in using different generator than StdGen.)
-}
module MagicHaskeller.MyCheck where
import System.Random
import Control.Monad(liftM, liftM2, liftM3)
import Data.Char(ord,chr)
import Data.Ratio

newtype Gen a = Gen {unGen :: Int -> StdGen -> a}
type Coarb a b = a -> Gen b -> Gen b

sized :: (Int -> Gen a) -> Gen a
sized fgen = Gen $ \n g -> unGen (fgen n) n g

instance Functor Gen where
    fmap = liftM

instance Monad Gen where
    return a    = Gen $ \_ _ -> a
    Gen m >>= k = Gen $ \n g -> case split g of (g1,g2) -> unGen (k (m n g1)) n g2

arbitraryR :: Random a => (a, a) -> Gen a
arbitraryR bnds = Gen $ \ _ gen -> fst $ randomR bnds gen
-- arbitrary :: (Random a, Bounded a) => Gen a
-- arbitrary = arbitraryR (minBound, maxBound)

arbitraryVoid :: Gen ()
arbitraryVoid = return ()
coarbitraryVoid :: Coarb () b
coarbitraryVoid _ = id

arbitraryBool :: Gen Bool
arbitraryBool = arbitraryR (False,True)
coarbitraryBool :: Coarb Bool b
-- coarbitraryBool b = if b then variant 0 else variant 1
coarbitraryBool b (Gen f) = Gen $ \size stdgen -> f size $ case split stdgen of (g0,g1) -> if b then g0 else g1

arbitraryInt :: Gen Int
arbitraryInt = arbitraryIntegral
coarbitraryInt :: Coarb Int b
coarbitraryInt n = newvariant n

arbitraryInteger :: Gen Integer
arbitraryInteger = arbitraryIntegral
coarbitraryInteger :: Coarb Integer b
coarbitraryInteger n = newvariant n

arbitraryIntegral :: (Random i, Integral i) => Gen i
arbitraryIntegral = sized $ \n -> arbitraryR ( - fromIntegral n,  fromIntegral n )

-- variant of Test.QuickCheck.variant using divide-and-conquer
logvariant, newvariant :: Integral i => i -> Gen a -> Gen a
logvariant 0 = coarbitraryBool True
logvariant n | n > 0 = coarbitraryBool False . logvariant (n `div` 2) . coarbitraryBool (n `mod` 2 == 0)
             | otherwise = error "logvariant: negative argument"
newvariant n | n >= 0    = coarbitraryBool True  . logvariant n
             | otherwise = coarbitraryBool False . logvariant (-1-n)

arbitraryFloat :: Gen Float
arbitraryFloat = arbitraryRealFloat
arbitraryDouble :: Gen Double
arbitraryDouble = arbitraryRealFloat

coarbitraryFloat :: Coarb Float b
coarbitraryFloat = coarbitraryRealFloat
coarbitraryDouble :: Coarb Double b
coarbitraryDouble = coarbitraryRealFloat

fraction a b c = fromInteger a + (fromInteger b / (abs (fromInteger c) + 1))

arbitraryRealFloat :: RealFloat a => Gen a
arbitraryRealFloat     = liftM3 fraction arbitraryInteger arbitraryInteger arbitraryInteger
coarbitraryRealFloat :: RealFloat a => Coarb a b
coarbitraryRealFloat x = let (sig, xpo) = decodeFloat x in newvariant sig . newvariant xpo

arbitraryChar = arbitraryR (' ', chr 126)
coarbitraryChar c = newvariant (ord c)

arbitraryOrdering :: Gen Ordering
arbitraryOrdering  = arbitraryR (0,2) >>= return . toEnum
-- Ordering is not an instance of Random!

arbitraryMaybe    :: Gen a -> Gen (Maybe a)
arbitraryMaybe arb = do b <- arbitraryBool
                        if b then return Nothing else liftM Just arb

arbitraryList     :: Gen a -> Gen [a]
arbitraryList  arb = sized $ \n -> arbitraryR (0,n) >>= \n -> sequence $ replicate n arb

arbitraryPair     :: Gen a -> Gen b -> Gen (a,b)
arbitraryPair      = liftM2 (,)

arbitraryEither   :: Gen a -> Gen b -> Gen (Either a b)
arbitraryEither arb0 arb1 = do b <- arbitraryBool
                               if b then liftM Left arb0 else liftM Right arb1

arbitraryTriplet  :: Gen a -> Gen b -> Gen c -> Gen (a,b,c)
arbitraryTriplet   = liftM3 (,,)

arbitraryFun :: Coarb a b -> Gen b -> Gen (a->b)
arbitraryFun coarb arb = Gen (\n r a -> unGen (coarb a arb) n r)

arbitraryRational :: Gen Rational
arbitraryRational = arbitrary


coarbitraryOrdering :: Coarb Ordering b
coarbitraryOrdering x = case x of LT -> coarbitraryBool True
                                  EQ -> coarbitraryBool False . coarbitraryBool True
                                  GT -> coarbitraryBool False . coarbitraryBool False

coarbitraryList :: Coarb a b -> Coarb [a] b
coarbitraryList _     []     = coarbitraryBool True
coarbitraryList coarb (x:xs) = coarbitraryBool False . coarb x . coarbitraryList coarb xs

coarbitraryMaybe :: Coarb a b -> Coarb (Maybe a) b
coarbitraryMaybe _     Nothing  = coarbitraryBool True
coarbitraryMaybe coarb (Just x) = coarbitraryBool False . coarb x

coarbitraryEither :: Coarb a c -> Coarb b c -> Coarb (Either a b) c
coarbitraryEither coarb0 _      (Left x)  = coarbitraryBool True . coarb0 x
coarbitraryEither _      coarb1 (Right y) = coarbitraryBool False . coarb1 y

coarbitraryPair :: Coarb a c -> Coarb b c -> Coarb (a,b) c
coarbitraryPair coarb0 coarb1 (a,b) = coarb0 a . coarb1 b

coarbitraryTriplet :: Coarb a d -> Coarb b d -> Coarb c d -> Coarb (a,b,c) d
coarbitraryTriplet coarb0 coarb1 coarb2 (a,b,c) = coarb0 a . coarb1 b . coarb2 c

coarbitraryFun :: Gen a -> Coarb b d -> Coarb (a->b) d
-- This is based on QuickCheck-1, and quite lightweight.
coarbitraryFun arb coarb f gen = arb >>= \x -> coarb (f x) gen

-- This is a definition based on QuickCheck-2:
-- coarbitraryFun arb coarb f gen = arbitraryList arb >>= \xs -> coarbitraryList coarb (map f xs) gen

-- This does even heavier check.
-- coarbitraryFun arb coarb f gen = (sized $ \n -> sequence $ replicate n arb) >>= \xs -> coarbitraryList coarb (map f xs) gen

class Arbitrary a where
    arbitrary   :: Gen a
class Coarbitrary a where
    coarbitrary :: a -> Gen b -> Gen b

instance Arbitrary () where
    arbitrary = arbitraryVoid
instance Coarbitrary () where
    coarbitrary = coarbitraryVoid

instance Arbitrary Bool where
    arbitrary = arbitraryBool
instance Coarbitrary Bool where
    coarbitrary = coarbitraryBool

instance Arbitrary Int where
    arbitrary = arbitraryInt
instance Coarbitrary Int where
    coarbitrary = coarbitraryInt

instance Arbitrary Integer where
    arbitrary = arbitraryInteger
instance Coarbitrary Integer where
    coarbitrary = coarbitraryInteger

instance Arbitrary Float where
    arbitrary = arbitraryFloat
instance Coarbitrary Float where
    coarbitrary = coarbitraryFloat

instance Arbitrary Double where
    arbitrary = arbitraryDouble
instance Coarbitrary Double where
    coarbitrary = coarbitraryDouble

instance Arbitrary Char where
    arbitrary = arbitraryChar
instance Coarbitrary Char where
    coarbitrary = coarbitraryChar

instance Arbitrary Ordering where
    arbitrary = arbitraryOrdering
instance Coarbitrary Ordering where
    coarbitrary = coarbitraryOrdering

instance Arbitrary a => Arbitrary (Maybe a) where
    arbitrary = arbitraryMaybe arbitrary
instance Coarbitrary a => Coarbitrary (Maybe a) where
    coarbitrary = coarbitraryMaybe coarbitrary

instance Arbitrary a => Arbitrary [a] where
    arbitrary = arbitraryList arbitrary
instance Coarbitrary a => Coarbitrary [a] where
    coarbitrary = coarbitraryList coarbitrary

instance (Arbitrary a, Arbitrary b) => Arbitrary (a,b) where
    arbitrary = arbitraryPair arbitrary arbitrary
instance (Coarbitrary a, Coarbitrary b) => Coarbitrary (a,b) where
    coarbitrary = coarbitraryPair coarbitrary coarbitrary

instance (Arbitrary a, Arbitrary b) => Arbitrary (Either a b) where
    arbitrary = arbitraryEither arbitrary arbitrary
instance (Coarbitrary a, Coarbitrary b) => Coarbitrary (Either a b) where
    coarbitrary = coarbitraryEither coarbitrary coarbitrary

instance (Arbitrary a, Arbitrary b, Arbitrary c) => Arbitrary (a,b,c) where
    arbitrary = arbitraryTriplet arbitrary arbitrary arbitrary
instance (Coarbitrary a, Coarbitrary b, Coarbitrary c) => Coarbitrary (a,b,c) where
    coarbitrary = coarbitraryTriplet coarbitrary coarbitrary coarbitrary

instance  (Coarbitrary a, Arbitrary b) => Arbitrary (a->b) where
    arbitrary = arbitraryFun coarbitrary arbitrary
instance  (Arbitrary a, Coarbitrary b) => Coarbitrary (a->b) where
    coarbitrary = coarbitraryFun arbitrary coarbitrary

instance (Integral i, Random i) => Arbitrary (Ratio i) where
    arbitrary  = liftM2 (%) arbitraryIntegral (fmap (\x->1+abs x) arbitraryIntegral)
instance (Integral i) => Coarbitrary (Ratio i) where
    coarbitrary r = newvariant (numerator r) . logvariant (denominator r)