{-# LANGUAGE FlexibleInstances #-}
{-# OPTIONS_GHC -Wno-orphans #-}
module Data.BenchRefinablePartition (benchmarks) where

import           Criterion

import           Control.Monad.ST
import           Data.Ord (comparing)

import           Control.DeepSeq

import           Data.RefinablePartition

benchmarks :: Benchmark
benchmarks = bgroup "Data.RefinablePartition"
                    [benchGroupBy, benchMake, benchMake1, benchMark]

benchGroupBy :: Benchmark
benchGroupBy = bgroup
  "groupBy"
  [ bench "10 states, already sorted"
    $ withInit (make 10 1 (const 0)) (\p -> groupBy p 0 (comparing id))
  , bench "20 states, already sorted"
    $ withInit (make 20 1 (const 0)) (\p -> groupBy p 0 (comparing id))
  , bench "100 states, already sorted"
    $ withInit (make 100 1 (const 0)) (\p -> groupBy p 0 (comparing id))
  , bench "1000 states, already sorted"
    $ withInit (make 1000 1 (const 0)) (\p -> groupBy p 0 (comparing id))
  , bench "10 states, reversed"
    $ withInit (make 10 1 (const 0)) (\p -> groupBy p 0 (comparing negate))
  , bench "20 states, reversed"
    $ withInit (make 20 1 (const 0)) (\p -> groupBy p 0 (comparing negate))
  , bench "100 states, reversed"
    $ withInit (make 100 1 (const 0)) (\p -> groupBy p 0 (comparing negate))
  , bench "1000 states, reversed"
    $ withInit (make 1000 1 (const 0)) (\p -> groupBy p 0 (comparing negate))
  ]


benchMake :: Benchmark
benchMake = bgroup
  "make"
  [ bench "1 block, 10 states" $ whnfIO (stToIO (make 10 1 (const 0)))
  , bench "1 block, 20 states" $ whnfIO (stToIO (make 20 1 (const 0)))
  , bench "1 block, 100 states" $ whnfIO (stToIO (make 100 1 (const 0)))
  , bench "2 blocks, 10 states" $ whnfIO (stToIO (make 10 2 (bmod 2)))
  , bench "2 blocks, 20 states" $ whnfIO (stToIO (make 20 2 (bmod 2)))
  , bench "2 blocks, 100 states" $ whnfIO (stToIO (make 100 2 (bmod 2)))
  , bench "10 blocks, 10 states" $ whnfIO (stToIO (make 10 10 Block))
  , bench "20 blocks, 20 states" $ whnfIO (stToIO (make 20 20 Block))
  , bench "100 blocks, 100 states" $ whnfIO (stToIO (make 100 100 Block))
  ]
  where bmod n = Block . (`mod` n)


benchMake1 :: Benchmark
benchMake1 = bgroup
  "make1"
  [ bench "10 states" $ whnfIO (stToIO (make1 10))
  , bench "20 states" $ whnfIO (stToIO (make1 20))
  , bench "100 states" $ whnfIO (stToIO (make1 100))
  ]


-- hehe
benchMark :: Benchmark
benchMark = bgroup
  "mark"
  [ bench "0 states" $ withInit (make1 100) (\_ -> return ())
  , bench "10 states" $ withInit (make1 100) (\p -> mapM_ (mark p) [0 .. 9])
  , bench "20 states" $ withInit (make1 100) (\p -> mapM_ (mark p) [0 .. 19])
  , bench "100 states" $ withInit (make1 100) (\p -> mapM_ (mark p) [0 .. 99])
  , bench "10 states, reverse"
    $ withInit (make1 100) (\p -> mapM_ (mark p) [9, 8 .. 0])
  , bench "20 states, reverse"
    $ withInit (make1 100) (\p -> mapM_ (mark p) [19, 18 .. 0])
  , bench "100 states, reverse"
    $ withInit (make1 100) (\p -> mapM_ (mark p) [99, 98 .. 0])
  ]

instance NFData (RefinablePartition RealWorld) where
  rnf p = seq p ()

withInit
  :: NFData a
  => ST RealWorld (RefinablePartition RealWorld)
  -> (RefinablePartition RealWorld -> ST RealWorld a)
  -> Benchmarkable
withInit initialize action = perRunEnv (stToIO initialize) (stToIO . action)