module Arbitrary3 where

import RandomNumbers
import Monad
import List
import Maybe

--============================================================================
-- We define a monad providing random numbers and a global size parameter.
--   choose low hi:  generates a random number in the given range
--   recursive f:    passes f the global size, for use in defining
--                   recursive generators.
--   bigger f:       run f with a bigger size
--============================================================================

newtype Generator a = G (Int -> Rand -> a)

runGenerator sz (G f) =
  do rg<-newRand
     return (f sz rg)

instance Monad Generator where
  return x = G (\sz rg -> x)
  G f >>= g = G$ \sz rg -> 
                 let (rg1,rg2) = split rg
                     G h = g (f sz rg1) in h sz rg2

instance Functor Generator where
  fmap f x = x >>= return . f

choose l h = G$ \sz rg -> l + fst (next rg) `mod` (h-l+1)

recursive f = G$ \sz rg -> let G g = f sz in g sz rg

bigger (G f) = G$ \sz rg-> f (sz+1) rg

promote :: (a -> Generator b) -> Generator (a->b)
promote f = G$ \sz rg -> \a -> let G g = f a in g sz rg

variant i (G g) = G$ \sz rg -> 
  g sz (iterate (snd.split) rg !! (i+1))

suchThat (G g) p = G$ satisfy
  where satisfy sz rg =
          let (rg1,rg2) = split rg
              x = g sz rg1
          in if p x then x else satisfy (sz+1) rg2

--============================================================================
-- We define a class for generation of arbitrary values, and some useful
-- instances.
--============================================================================

class Arbitrary a where
  arbitrary :: Generator a

instance Arbitrary Int where
  arbitrary = recursive$ \sz -> choose (-sz) sz

instance Arbitrary Integer where
  arbitrary = liftM fromInt arbitrary

instance Arbitrary Char where
  arbitrary = do c<-choose 0 127
                 if isPrint (chr c) then return (chr c) else arbitrary

instance Arbitrary Double where
  arbitrary = let big = 10^8 in
              do a<-choose (-big) big
                 b<-choose 0 big
                 c<-choose (-28) 12
                 return$ (fromInt a + fromInt b / fromInt big) * 10^^c

instance Arbitrary Float where
  arbitrary = fmap fromDouble arbitrary

instance Arbitrary () where
  arbitrary = return ()

instance (Arbitrary a, Arbitrary b) => Arbitrary (a,b) where
  arbitrary = liftM2 (,) arbitrary arbitrary

instance Arbitrary a => Arbitrary (Maybe a) where
  arbitrary = arbitraryMaybe arbitrary

arbitraryMaybe a = oneofWithFreq [(1, return Nothing),
                                  (2, liftM Just a)]

instance Arbitrary a => Arbitrary [a] where
  arbitrary = arbitraryList arbitrary

arbitraryList a = recursive arbList
    where arbList 0 = return []
          arbList n = oneofWithFreq [(1,return []), 
                                     (4,liftM2 (:) a (arbList (n-1)))]

arbitraryElement xs = do i<-choose 0 (length xs - 1)
                         return (xs!!i)

--============================================================================
-- The following combinators are intended for use in defining generators
-- for datatypes.
--============================================================================

oneof gs = oneofWithFreq (zip [1,1..] gs)

oneofWithFreq gs = do c<-choose 1 (sum (map fst gs))
                      find c gs
  where find c ((n,g):gs) =
          if c<=n then g else find (c-n) gs

--============================================================================
-- Generation of arbitrary functions
--============================================================================

arbitraryFn coa arb = promote (\a -> coa a arb)

instance (CoArbitrary a, Arbitrary b) => Arbitrary (a->b) where
  arbitrary = arbitraryFn coarbitrary arbitrary
  
class CoArbitrary a where
  coarbitrary :: a -> Generator b -> Generator b

instance CoArbitrary Bool where
  coarbitrary b = if b then variant 0 else variant 1

instance CoArbitrary Integer where
  coarbitrary n = coarbitrary (bits (if n<0 then (-2*n)-1 else 2*n))
    where bits 0 = []
          bits n = even n:bits (n`div`2)

instance CoArbitrary Int where
  coarbitrary n = coarbitrary (toInteger n)

instance CoArbitrary Char where
  coarbitrary c = coarbitrary (ord c)

instance CoArbitrary Double where
  coarbitrary d = coarbitrary (decodeFloat d)

instance CoArbitrary Float where 
  coarbitrary f = coarbitrary (decodeFloat f)

instance CoArbitrary () where
  coarbitrary () = id

instance (CoArbitrary a, CoArbitrary b) => CoArbitrary (a,b) where
  coarbitrary (a,b) = coarbitrary a . coarbitrary b

instance CoArbitrary a => CoArbitrary (Maybe a) where
  coarbitrary Nothing = variant 0
  coarbitrary (Just a) = variant 1 . coarbitrary a

instance CoArbitrary a => CoArbitrary [a] where
  coarbitrary [] = variant 0
  coarbitrary (x:xs) = variant 1 . coarbitrary xs . coarbitrary x

coarbitraryElement xs x =
  coarbitrary (fromJust (x `elemIndex` xs))

instance (Arbitrary a, CoArbitrary b) => CoArbitrary (a->b) where
  coarbitrary f g = 
    do n<-arbitrary
       let m = round(sqrt(fromInt (abs n)))
       as<-sequence [arbitrary | i<-[0..m]]
       foldr id g (map (coarbitrary.f) as)

--============================================================================
-- Functions to test Generators
--============================================================================

testGenerator g =
  do xs<-runGenerator 20 (sequence [g | i<-[1..20]])
     putStr (unlines (map show xs))