{-# LANGUAGE ScopedTypeVariables #-}
import System.Random
import Data.Word
import Data.Bits
import IO
import Data.Array
import Data.Ratio
import System

randomR_ranged :: forall a g. (RandomGen g, Integral a) => (a,a) -> g -> (a,g)
randomR_ranged (lo,hi) gen = let
    integerRandom :: (Integer,g)
    integerRandom = randomR (fromIntegral lo,fromIntegral hi) gen
    in (fromInteger (fst integerRandom),snd integerRandom)

instance Random Word8 where
    randomR = randomR_ranged
    random = randomR (minBound, maxBound)

instance Random Word16 where
    random = randomR (minBound, maxBound)
    randomR = randomR_ranged

{-
  For the Word32 and Word64 "random" functions
  (returning a random number in the full range),
  we can use a more optimized verison that calls
  random :: Word16 multiple times and concatenates
  the bits.  This avoids some of the overhead of
  gmp and Integer arithmetic.
-}

paste_two :: forall g small big. 
   (RandomGen g, Num big, Integral small, Bits big, Bits small) => 
   (g -> (small, g)) -> g -> (big, g)
paste_two small_generator gen = let
    lower :: (small,g)
    lower = small_generator gen
    lowerbits :: big
    lowerbits = fromIntegral (fst lower)
    upper :: (small,g)
    upper = small_generator (snd lower)
    upperbits :: big
    upperbits = fromIntegral (fst upper)
    in ((shiftL upperbits (bitSize (fst lower))) .|. lowerbits,
        (snd upper))

random_Word32 :: forall g. RandomGen g => g -> (Word32, g)
random_Word32 gen = let
    small_generator :: g -> (Word16, g)
    small_generator = random
    in paste_two small_generator gen

random_Word64 :: forall g. RandomGen g => g -> (Word64, g)
random_Word64 gen = let
    small_generator :: g -> (Word32, g)
    small_generator = random
    in paste_two small_generator gen

instance Random Word32 where
    random = random_Word32
    randomR = randomR_ranged

instance Random Word64 where
    random = random_Word64
    randomR = randomR_ranged

testByte :: IO ()
testByte =
  sequence_ (replicate 25600000 
             (do
               (w::Word8) <- randomIO
               print w
             ))
{-
this results in chi-squared=
274.15626000000000000000
and p=0.1951
-}

testShorthi :: IO ()
testShorthi =
    sequence_ (replicate 25600000
                          (do
                            (w::Word16) <- randomIO
                            print (shiftR w 8)
                          ))
  -- chiSQ=281.56036

testShortlo= do
  hPutStrLn stderr "testShortlo"
  sequence_ (replicate 2560000
                           (do
                             (w::Word16) <- randomIO
                             print (w .&. 255)
                           ))
  --267.1608

countBytes :: [Word8] -> [Int]
countBytes bs = elems$accumArray (+) 0 (minBound,maxBound) (zip bs (repeat 1))

chisq :: Int -> [Int] -> Rational
chisq expected observed =
    let
        diff :: Int -> Integer
        diff x = (toInteger (x-expected))^2
        alldiffs :: Integer
        alldiffs = sum$map diff observed
    in
      alldiffs % (toInteger expected)

run_chisq :: Int -> (IO Word8) -> IO Double
run_chisq count f = do
  (o :: [Word8]) <- sequence (replicate (256*count) f)
  return (fromRational(chisq count (countBytes o)))

word16shift :: Int -> Word16 -> Word8
word16shift i x = fromIntegral(shiftR x i)

word32shift :: Int -> Word32 -> Word8
word32shift i x = fromIntegral(shiftR x i)

word64shift :: Int -> Word64 -> Word8
word64shift i x = fromIntegral(shiftR x i)

count = 1000

main :: IO ()
main = do
  args <- getArgs
  setStdGen(read (head args))
  run_chisq count (randomIO) >>= print
  run_chisq count (randomIO >>= (return . (word16shift 0))) >>= print
  run_chisq count (randomIO >>= (return . (word16shift 8))) >>= print
  putStrLn ""
  run_chisq count (randomIO >>= (return . (word32shift 0))) >>= print
  run_chisq count (randomIO >>= (return . (word32shift 8))) >>= print
  run_chisq count (randomIO >>= (return . (word32shift 16))) >>= print
  run_chisq count (randomIO >>= (return . (word32shift 24))) >>= print
  putStrLn ""
  run_chisq count (randomIO >>= (return . (word64shift 0))) >>= print
  run_chisq count (randomIO >>= (return . (word64shift 8))) >>= print
  run_chisq count (randomIO >>= (return . (word64shift 16))) >>= print
  run_chisq count (randomIO >>= (return . (word64shift 24))) >>= print
  run_chisq count (randomIO >>= (return . (word64shift 32))) >>= print
  run_chisq count (randomIO >>= (return . (word64shift 40))) >>= print
  run_chisq count (randomIO >>= (return . (word64shift 48))) >>= print
  run_chisq count (randomIO >>= (return . (word64shift 56))) >>= print