{-# OPTIONS_GHC -cpp -fglasgow-exts -fallow-overlapping-instances #-}
import Control.Concurrent
import Control.Exception
import Control.Monad
import Data.List
import Data.Word
import Foreign.C.String
import Foreign.Marshal.Alloc
import Foreign.Storable
import Numeric
import System.CPUTime
import System.Directory
import System.Environment
import System.Exit
import System.IO
import System.Mem
import System.Time

import Data.Array.Base
import Data.Array.IO
import Data.Array.Unboxed
import Binary

size  = 10^8 :: Int    -- whole size of data processed in each test
chunk = 10^4 :: Int    -- size of chunk what is read/written in each operation
type Element = Word32  -- type of elements used for testing

-- Automatically calculated
bufsize = fromIntegral (size+times*10)              -- size of buffer allocated
times = size `div` chunk                            -- count of repetitions
elements = chunk `div` sizeOf (undefined::Element)  -- number of elements in array/list
tempname = "BinaryBenchmark.tmp"                    -- name of temporary file used in test

main = do
    putStrLn ("Binary libraries benchmark."++
              " Using "++showM size++" of data in chunks of "++showM chunk)
    try (removeFile tempname)

    flip finally (try (removeFile tempname)) $ do
    test "GHC 6.5 Binary Memory" $ openBinMem bufsize
    test "GHC 6.5 Binary File"   $ (openBinaryFile tempname ReadWriteMode >>= openBinIO )

{-# INLINE test #-}
test name open = do
    putStrLn ""
    putStrLn (name++":")
    h <- open
    o <- tellBin h

    let arr = array (0,elements-1) [] :: UArray Int Element
    return $! (arr!0 )  -- to make the following timing accurate
    benchmark h "Writing UArray" times (putUArray h arr)
    performGC
    seekBin h o
    benchmark h "Reading UArray" times (getUArray h :: IO (IOUArray Int Element))
    performGC

    let list = replicate elements (0::Element)
    return $! (length list)  -- to make the following timing accurate
    seekBin h o
    benchmark h "Writing list" times (putNListJhc h list)
    performGC
    seekBin h o
    benchmark h "Reading list" times (getNListJhc h :: IO [Element])
    performGC


{-# INLINE putNListJhc #-}
-- Put list
putNListJhc :: Binary a => BinHandle -> [a] -> IO ()
putNListJhc bh xs = do
    put_ bh (length xs)
    mapM_ (put_ bh) xs

{-# INLINE getNListJhc #-}
-- Get list
getNListJhc :: Binary a => BinHandle -> IO [a]
getNListJhc bh = do
    n <- get bh
    sequence $ replicate n (get bh)

{-# INLINE putUArray #-}
-- Put UArray
putUArray h arr = do put_ h (bounds arr)
                     mapRange 0 (rangeSize(bounds arr)-1)
                         (put_ h . unsafeAt arr)

{-# INLINE getUArray #-}
-- Get UArray
getUArray h = do bounds <- get h
                 arr <- newArray_ bounds :: IO (IOUArray Int Element)
                 mapRange 0 (rangeSize(bounds)-1)
                     (\i -> get h >>= unsafeWrite arr i)
                 return arr

{-# INLINE mapRange #-}
-- Faster equivalent of "mapM_ action [from..to]"
mapRange from to action  = go from
  where
    go i | i>to      = return ()
         | otherwise = do action i
                          go $! (i+1)


{-# INLINE benchmark #-}
benchmark h str times action = do
  handle (\_ -> putStrLn$ str ++ ": failed!") $ do
      prev <- getCPUTime
      prev2 <- getClockTime
      loop times action
      current <- getCPUTime
      current2 <- getClockTime
      let secs = fromIntegral (current-prev) / 1e12
          secs2 = diffTimes current2 prev2
      putStrLn$ str ++ ": " ++ showTime secs2 ++ " (user: " ++ showTime secs ++ ")"

showTime secs  =  showFFloat (Just 3) secs " secs"

diffTimes (TOD sa pa) (TOD sb pb)  =  i(sa - sb) + (i(pa-pb) / 1e12)

i x = fromIntegral x

{-# INLINE loop #-}
loop n action  = go n
  where
    go n | n `seq` False = undefined
    go 0 = return ()
    go n = do action
              go (n-1)

showM x | x `mod` 10^9 == 0   = show (x `div` 10^9) ++ "gb"
        | x `mod` 10^6 == 0   = show (x `div` 10^6) ++ "mb"
        | x `mod` 10^3 == 0   = show (x `div` 10^3) ++ "kb"
        | otherwise           = show x              ++ " bytes"