import "Prelude"
import "Random"
import "IterN"
import "Vector" as Vector

data Gen a = Gen (Integer -> <Random> a)

@inline
genRandom :: (<Random> a) -> Gen a
genRandom r = Gen (\_ -> r)

@inline
runGen :: Gen a -> Integer -> <Random> a
runGen (Gen f) n = f n

class Arbitrary a where
    arbitrary :: Gen a

class CoArbitrary a where
    variateSeed :: Integer -> Long -> a -> Long

instance Arbitrary Boolean where
    arbitrary = genRandom randomBoolean
instance CoArbitrary Boolean where
    variateSeed _ seed v = if v then seed + 1 else seed

@private
@inline
negateRandomly :: Integer -> <Random> Integer
negateRandomly v = if randomBoolean then v else -v

instance Arbitrary Integer where
    arbitrary = Gen $ \n -> negateRandomly $ randomN $ n+1
instance CoArbitrary Integer where
    variateSeed _ seed v = seed + fromInteger v
    
instance Arbitrary Long where
    arbitrary = genRandom randomLong
instance CoArbitrary Long where
    variateSeed _ seed v = seed + v
    
instance Arbitrary Float where
    arbitrary = Gen $ \n -> fromInteger n 
                          * (fromInteger $ negateRandomly $ randomN $ pres)
                          / (fromInteger $ 1+randomN pres)
        where pres = 999999999
instance CoArbitrary Float where
    variateSeed _ seed v = seed + fromInteger (floatToIntBits v)

instance Arbitrary Double where
    arbitrary = Gen $ \n -> fromInteger n 
                          * (fromInteger $ negateRandomly $ randomN $ pres)
                          / (fromInteger $ 1+randomN pres)
        where pres = 999999999
instance CoArbitrary Double where
    variateSeed _ seed v = seed + doubleToLongBits v

randomCharacter :: <Random> Character
randomCharacter = let p = randomN 100 
                  in if p < 80
                     then addChar ' ' (randomN 95)
                     else if p < 90
                     then addChar '\0' (randomN 256)
                     else addChar '\0' (randomN 65536)
randomIdentifierFirstCharacter :: <Random> Character
randomIdentifierFirstCharacter = if randomBoolean
                                 then addChar 'a' (randomN 26)
                                 else addChar 'A' (randomN 26)
randomIdentifierCharacter :: <Random> Character
randomIdentifierCharacter = let p = randomN 100 
                            in if p < 70
                               then addChar 'a' (randomN 26)
                               else if p < 90
                               then addChar 'A' (randomN 26)
                               else addChar '0' (randomN 10)
                        
instance Arbitrary Character where
    arbitrary = genRandom randomCharacter
instance CoArbitrary Character where
    variateSeed _ seed v = seed + fromInteger (hashCode (showCharacter v))

"""Generates a random string that is a valid SCL or Java identifier""" 
arbitraryIdentifier = Gen $ \n -> string (Vector.vectorF (1+randomN (max n 1))
    (\i -> if i==0 then randomIdentifierFirstCharacter else randomIdentifierCharacter))

instance Arbitrary String where
    arbitrary = Gen $ \n -> string (Vector.vectorF (randomN (1+n)) (\_ -> randomCharacter))
instance CoArbitrary String where
    variateSeed _ seed v = seed + fromInteger (hashCode v)

instance Arbitrary () where
    arbitrary = genRandom ()
instance CoArbitrary () where
    variateSeed _ seed _ = seed

instance (Arbitrary a, Arbitrary b) => Arbitrary (a, b) where
    arbitrary = arbitraryTuple2 arbitrary arbitrary
instance (CoArbitrary a, CoArbitrary b) => CoArbitrary (a, b) where
    variateSeed n seed (a, b) = variateSeed n (13*variateSeed n seed a) b

instance (Arbitrary a, Arbitrary b, Arbitrary c) => Arbitrary (a, b, c) where
    arbitrary = arbitraryTuple3 arbitrary arbitrary arbitrary
instance (CoArbitrary a, CoArbitrary b, CoArbitrary c) => CoArbitrary (a, b, c) where
    variateSeed n seed (a, b, c) = variateSeed n (13*variateSeed n (13*variateSeed n seed a) b) c

instance (Arbitrary a, Arbitrary b, Arbitrary c, Arbitrary d) => Arbitrary (a, b, c, d) where
    arbitrary = arbitraryTuple4 arbitrary arbitrary arbitrary arbitrary
instance (CoArbitrary a, CoArbitrary b, CoArbitrary c, CoArbitrary d) => CoArbitrary (a, b, c, d) where
    variateSeed n seed (a, b, c, d) = variateSeed n (13*variateSeed n (13*variateSeed n (13*variateSeed n seed a) b) c) d

instance (Arbitrary a, Arbitrary b, Arbitrary c, Arbitrary d, Arbitrary e) => Arbitrary (a, b, c, d, e) where
    arbitrary = arbitraryTuple5 arbitrary arbitrary arbitrary arbitrary arbitrary
instance (CoArbitrary a, CoArbitrary b, CoArbitrary c, CoArbitrary d, CoArbitrary e) 
      => CoArbitrary (a, b, c, d, e) where
    variateSeed n seed (a, b, c, d, e) = 
        variateSeed n (13*variateSeed n (13*variateSeed n (13*variateSeed n (13*variateSeed n seed a) b) c) d) e

instance (Arbitrary a, Arbitrary b, Arbitrary c, Arbitrary d, Arbitrary e, Arbitrary f) 
      => Arbitrary (a, b, c, d, e, f) where
    arbitrary = arbitraryTuple6 arbitrary arbitrary arbitrary arbitrary arbitrary arbitrary
instance (CoArbitrary a, CoArbitrary b, CoArbitrary c, CoArbitrary d, CoArbitrary e, CoArbitrary f) 
      => CoArbitrary (a, b, c, d, e, f) where
    variateSeed n seed (a, b, c, d, e, f) = 
        variateSeed n (13*variateSeed n (13*variateSeed n (13*variateSeed n (13*variateSeed n (13*variateSeed n seed a) b) c) d) e) f

instance (Arbitrary a) => Arbitrary (Maybe a) where
    arbitrary = arbitraryMaybe arbitrary
instance (CoArbitrary a) => CoArbitrary (Maybe a) where
    variateSeed n seed v = match v with
        Nothing -> 13*seed+1
        Just a -> 13*variateSeed n seed a

instance (Arbitrary a, Arbitrary b) => Arbitrary (Either a b) where
    arbitrary = arbitraryEither arbitrary arbitrary
instance (CoArbitrary a, CoArbitrary b) => CoArbitrary (Either a b) where
    variateSeed n seed v = match v with
        Left a  -> 13*variateSeed n seed a
        Right a -> 13*variateSeed n seed a+1

instance (Arbitrary a) => Arbitrary [a] where
    arbitrary = arbitraryList arbitrary
instance (CoArbitrary a) => CoArbitrary [a] where
    variateSeed n = foldl (variateSeed n)
        
instance (Arbitrary a, VecComp a) => Arbitrary (Vector a) where
    arbitrary = arbitraryVector arbitrary
instance (CoArbitrary a) => CoArbitrary (Vector a) where
    variateSeed n = Vector.foldlVector (variateSeed n)
    
instance (CoArbitrary a, Arbitrary b) => Arbitrary (a -> b) where
    arbitrary = Gen $ \n -> let seed = randomLong
                            in  \v -> withSeed (variateSeed n seed v) (runGen arbitrary n)
instance (Arbitrary a, CoArbitrary b) => CoArbitrary (a -> b) where
    variateSeed n seed f = variateSeed n seed (withSeed seed (mapList f (runGen arbitrary n)))

arbitraryTuple2 (Gen a) (Gen b) = Gen $ \n -> (a n, b n)

arbitraryTuple3 (Gen a) (Gen b) (Gen c) = Gen $ \n -> (a n, b n, c n)

arbitraryTuple4 (Gen a) (Gen b) (Gen c) (Gen d) = Gen $ \n -> (a n, b n, c n, d n)

arbitraryTuple5 (Gen a) (Gen b) (Gen c) (Gen d) (Gen e) = Gen $ \n -> (a n, b n, c n, d n, e n)

arbitraryTuple6 (Gen a) (Gen b) (Gen c) (Gen d) (Gen e) (Gen f) = Gen $ \n -> (a n, b n, c n, d n, e n, f n)

arbitraryMaybe (Gen f) =
    Gen $ \n -> if randomN 4 < 1
                then Nothing
                else Just (f n)

arbitraryEither (Gen left) (Gen right) =
    Gen $ \n -> if randomBoolean
                then Left (left n)
                else Right (right n)

arbitraryFixedSizeList :: Integer -> Gen a -> Gen [a]
arbitraryFixedSizeList s (Gen f) = Gen $ \n -> mapN (\_ -> f n) s

arbitraryList :: Gen a -> Gen [a]
arbitraryList (Gen f) = Gen $ \n -> mapN (\_ -> f n) (randomN n)

@inline
arbitraryFixedSizeVector :: VecComp a => Integer -> Gen a -> Gen (Vector a)
arbitraryFixedSizeVector s (Gen f) = Gen $ \n -> Vector.vectorF s (\_ -> f n)

@inline
arbitraryVector :: VecComp a => Gen a -> Gen (Vector a)
arbitraryVector (Gen f) = Gen $ \n -> Vector.vectorF (randomN n) (\_ -> f n)

arbitraryElement :: [a] -> Gen a
arbitraryElement elements = genRandom (elements ! randomN (length elements))

arbitraryResized :: (Integer -> Integer) -> Gen a -> Gen a
arbitraryResized f (Gen gen) = Gen (gen . f)