import Criterion.Main
import Criterion.Config

import Data.Word
import Data.Bits
import System.Random

import qualified Data.Vector.Unboxed as UVec
import qualified Data.Vector.Storable as SVec
import qualified Data.Vector.Primitive as PVec

import qualified Data.ByteString as BStr
import qualified Data.ByteString.Unsafe as BStr

import qualified Data.Array.Unboxed as UArr
import qualified Data.Array.Storable as SArr

num = 1000000 :: Int
logNum = floor $ logBase 2 (fromIntegral num)
maskNum = 2^logNum-1

gen = mkStdGen 4653

randomList = take num $ randoms gen

-- PREPARE ARRAYS
uvec = UVec.fromList randomList :: UVec.Vector Word8
svec = SVec.fromList randomList :: SVec.Vector Word8
pvec = PVec.fromList randomList :: PVec.Vector Word8
bstr = BStr.pack     randomList :: BStr.ByteString
uarr = UArr.listArray (0,num-1) randomList :: UArr.UArray Int Word8

-- FOR SECOND TEST
randomAccessSum :: (Int -> Word8) -> Word8
randomAccessSum f = go (num-1) 0 0 where
    go 0 _ result = result
    go n index oldResult = let result = f index + oldResult in
        seq result go (n-1) ((index + 38634329) .&. maskNum) result

-- myConfig = defaultConfig { cfgVerbosity = ljust Quiet }
myConfig = defaultConfig

main = do
    uvec `seq` svec `seq` pvec `seq` bstr `seq` return ()
    defaultMainWith myConfig (return ())
        [ bgroup "sum"
            [ bench "Unboxed Vector: sum"           $ whnf UVec.sum uvec
            --, bench "Unboxed Vector: foldl1' (+)"   $ whnf (UVec.foldl1' (+)) uvec
            , bench "Storable Vector: sum"          $ whnf SVec.sum svec
            , bench "Storable Vector: foldl1' (+)"  $ whnf (SVec.foldl1' (+)) svec
            --, bench "Storable Vector: sum . toList" $ whnf (sum . SVec.toList) svec
            , bench "Primitive Vector: sum"         $ whnf PVec.sum pvec
            , bench "ByteString: foldl1' (+)"       $ whnf (BStr.foldl1' (+)) bstr
            , bench "ByteString: foldl1 (+)"      $ whnf (BStr.foldl1 (+)) bstr
            ]
        , bgroup "random access sum"
            [ bench "Unboxed Vector: unsafeIndex"   $ whnf randomAccessSum (UVec.unsafeIndex uvec)
            , bench "Storable Vector: unsafeIndex"  $ whnf randomAccessSum (SVec.unsafeIndex svec)
            , bench "Storable Vector: safe (!)"     $ whnf randomAccessSum ((SVec.!) svec)
            , bench "Primitive Vector: unsafeIndex" $ whnf randomAccessSum (PVec.unsafeIndex pvec)
            , bench "ByteString: unsafeIndex"       $ whnf randomAccessSum (BStr.unsafeIndex bstr)
            , bench "ByteString: safe index"        $ whnf randomAccessSum ((BStr.index) bstr)
            , bench "UArray: safe (!)"              $ whnf randomAccessSum ((UVec.!) uvec)
            ]
        ]