{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RankNTypes   #-}

module Criterion.Collection.Sample
  ( Benchmark(..)
  , SampleData(..)
  , MeasurementMode(..)
  , WorkloadMode(..)
  , computeMeanAndStddev
  , compute95thPercentile
  , computeMax
  , takeSample
  , takeSamples
  ) where

import           Control.DeepSeq
import           Control.Monad
import           Control.Monad.Trans
import           Criterion                           hiding (Benchmark)
import           Criterion.Collection.Internal.Types
import           Criterion.Config
import           Criterion.Environment
import           Criterion.IO.Printf
import           Criterion.Measurement
import           Criterion.Monad
import           Data.IORef
import           Data.List                           (foldl')
import           Data.Monoid
import qualified Data.Vector                         as V
import qualified Data.Vector.Unboxed                 as U
import           Statistics.Quantile                 (cadpw, continuousBy)
import           Statistics.Sample
import           System.Mem                          (performGC)
import           System.Random.MWC
import           Text.Printf                         (printf)

------------------------------------------------------------------------------
data MeasurementMode = PerBatch | PerOperation
data WorkloadMode = Pure | Mutating


------------------------------------------------------------------------------
data SampleData = SampleData {
      sdInputSize :: !Int       -- ^ what was the size of the input for this
                                -- sample?
    , sdNumOps    :: !Int       -- ^ how many operations are covered by this
                                -- sample? For a per-operation measurement,
                                -- this value would be \"1\", and for a batch
                                -- measurement this value would be the number
                                -- of items in the batch.
    , sdData      :: !Sample    -- ^ sample data.
    }


instance Show SampleData where
    show (SampleData is nop da) = "<SampleData inputSize=" ++ show is
                                  ++ ", nops=" ++ show nop
                                  ++ ", sample size=" ++ show (U.length da)
                                  ++ ">"

------------------------------------------------------------------------------
data Benchmark op = Benchmark {
      benchmarkName     :: String
    , dataStructures    :: [(String, DataStructure op)]
    , inputSizes        :: [Int]
    , workloadGenerator :: WorkloadGenerator op
}


------------------------------------------------------------------------------
-- | Given some sample data, compute the mean time per operation (in seconds)
-- and standard deviation
computeMeanAndStddev :: SampleData -> (Double, Double)
computeMeanAndStddev (SampleData _ nops sample) = (v,s)
  where
    nopsReal         = fromIntegral nops
    (meanValue, var) = meanVarianceUnb sample
    stddev           = sqrt $ abs var
    !v               = meanValue / nopsReal
    !s               = stddev / nopsReal


------------------------------------------------------------------------------
-- | Given some sample data, compute the 95th percentile.
compute95thPercentile :: SampleData -> Double
compute95thPercentile (SampleData _ nops sample) = v
  where
    nopsReal         = fromIntegral nops
    quantile         = continuousBy cadpw 19 20 sample
    v                = quantile / nopsReal


------------------------------------------------------------------------------
-- | Given some sample data, compute the maximum value
computeMax :: SampleData -> Double
computeMax (SampleData _ nops sample) = v
  where
    nopsReal         = fromIntegral nops
    maxval           = U.foldl' max 0 sample
    v                = maxval / nopsReal


------------------------------------------------------------------------------
takeSample :: (NFData op) =>
              MeasurementMode
           -> WorkloadMode
           -> Benchmark op
           -> Environment
           -> GenIO
           -> Int
           -> Criterion [SampleData]
takeSample !mMode !wMode !benchmark !env !rng !inputSize = do
    workload <- liftIO $ runWorkloadMonad (workGen inputSize) rng
    let setupOperations = setupWork workload
    let genWorkData     = genWorkload workload

    case mMode of
      PerBatch     -> batch setupOperations genWorkData
      PerOperation -> perOp setupOperations genWorkData


  where
    --------------------------------------------------------------------------
    dss       = dataStructures benchmark
    workGen   = workloadGenerator benchmark

    --------------------------------------------------------------------------
    batch setupOperations genWorkData = do
        workData <- liftIO $ runWorkloadMonad (genWorkData $ inputSize `div` 2)
                                              rng
        let nOps = V.length workData
        mapM (batchOne setupOperations workData nOps) dss


    --------------------------------------------------------------------------
    mkRunOp runOpMutating =
        let runOpPure = \m op -> do
                            m' <- runOpMutating m op
                            return $! m' `seq` m
        in case wMode of
             Pure     -> runOpPure
             Mutating -> runOpMutating


    --------------------------------------------------------------------------
    runWorkData workData chunkSize runOp start i val = go i val
      where
        go !i !val
            | i >= chunkSize = return val
            | otherwise = do
                  let op = V.unsafeIndex workData (start+i)
                  !val' <- runOp val op
                  go (i+1) val'


    --------------------------------------------------------------------------
    batchOne setupOperations workData nOps
             (name, (DataStructure emptyValue runOpMutating)) = do
        note $ "running batch benchmark on " ++ name ++ "\n"
        let minTime = envClockResolution env * 1000
        cfg <- getConfig
        let proc = V.foldM' runOpMutating
        let mkStartValue = emptyValue inputSize >>= flip proc setupOperations
        startValue1 <- liftIO mkStartValue

        liftIO performGC
        let tProc = runWorkData workData nOps runOpMutating 0 0

        prolix $ "running test batch with " ++ show nOps
                 ++ " work items\n"

        (tm,_) <- liftIO $ time (tProc startValue1)

        prolix $ "running initial timing on " ++ show nOps
                 ++ " work items\n"

        let iters = max 5 (ceiling (minTime / tm))

        prolix $ "running benchmark on " ++ show nOps
                 ++ " work items, " ++ show iters ++ " iterations\n"

        sample <- liftIO $ U.generateM iters $ \_ -> do
                      sv <- mkStartValue
                      performGC
                      (!tm,_) <- time (tProc sv)
                      return $ tm - clockCost

        prolix $ "finished batch benchmark on " ++ name ++ "\n"
        return (SampleData inputSize nOps sample)


    --------------------------------------------------------------------------
    perOp setupOperations genWorkData = do
        -- FIXME: lifted this code from criterion, is there some way to merge
        -- them?
        _ <- prolix "generating seed workload"
        seedData <- liftIO $ runWorkloadMonad (genWorkData 1000) rng
        _ <- prolix "seed workload generated"

        workData <- liftIO $ runWorkloadMonad (genWorkData inputSize) rng
        mapM (perOpOne setupOperations workData seedData) dss


    --------------------------------------------------------------------------
    perOpOne setupOperations workData seedData
             (name, (DataStructure emptyValue runOpMutating)) = do
        let runOp = mkRunOp runOpMutating
        let proc = V.foldM' runOpMutating

        note $ "running per-op benchmark on " ++ name ++ "\n"

        startValue <- liftIO (emptyValue inputSize >>=
                              flip proc setupOperations)
        liftIO performGC

        -- warm up clock
        _ <- liftIO $ runForAtLeast 0.1 10000 (`replicateM_` getTime)
        let minTime = envClockResolution env * 1000

        (testTime, testIters, startValue') <-
            liftIO $ timeSeed (min minTime 0.1) seedData runOp startValue
        _   <- note "ran %d iterations in %s\n" testIters (secs testTime)
        cfg <- getConfig

        let testItersD       = fromIntegral testIters
        let sampleCount      = fromLJ cfgSamples cfg
        let timePer          = (testTime - testItersD * clockCost) / testItersD
        let chunkSizeD       = minTime / timePer
        let chunkSize        = min (V.length workData) (ceiling chunkSizeD)
        let nSamples1        = min (chunkSize * sampleCount) (V.length workData)
        let numItersD        = fromIntegral nSamples1 / fromIntegral chunkSize
        let nSamples = max 1 (floor numItersD * chunkSize)

        _ <- note "collecting %d samples (in chunks of %d) in estimated %s\n"
                  nSamples chunkSize
                  (secs ((chunkSizeD * timePer + clockCost)*numItersD))


        (sample,_) <- mkSample chunkSize nSamples workData startValue' runOp
        liftIO performGC
        return (SampleData inputSize chunkSize sample)


    --------------------------------------------------------------------------
    mkSample chunkSize nSamples workData startValue runOp = liftIO $ do
        valRef <- newIORef startValue

        let numItersD = fromIntegral nSamples / fromIntegral chunkSize
        -- make sure nSamples is an integral multiple of chunkSize
        let numIters = max 1 (floor (numItersD :: Double))

        sample <- U.generateM numIters $ \chunk -> do
            !val <- readIORef valRef
            (!tm, val') <- time (runWorkData workData chunkSize runOp
                                             (chunk*chunkSize) 0 val)
            writeIORef valRef val'
            return $ tm - clockCost

        val <- readIORef valRef
        return (sample :: U.Vector Double, val)

    --------------------------------------------------------------------------
    clockCost = envClockCost env

    --------------------------------------------------------------------------
    timeSeed howLong seedData runOp startValue =
        loop startValue seedData (0::Int) =<< getTime
      where
        loop sv seed iters initTime = do
            now <- getTime
            let n = V.length seed
            when (now - initTime > howLong * 10) $
                 fail (printf "took too long to run: seed %d, iters %d"
                              (V.length seed) iters)
            (elapsed, (_,sv')) <- time (mkSample 1 n seed sv runOp)

            if elapsed < howLong
              then loop sv' (seed `mappend` seed) (iters+1) initTime
              else return (elapsed, n, sv')


------------------------------------------------------------------------------
takeSamples :: (NFData op) =>
               MeasurementMode
            -> WorkloadMode
            -> Benchmark op
            -> Environment
            -> GenIO
            -> Criterion [(String, [SampleData])]
takeSamples !mMode !wMode !benchmark !env !rng = do
    let szs = inputSizes benchmark
    when (null szs) $ fail "No input sizes defined"

    ssamples <- mapM (takeSample mMode wMode benchmark env rng) szs
    let names = map fst $ dataStructures benchmark
    let inputs = foldl' combine (map (const []) names) (reverse ssamples)

    return $ names `zip` inputs

  where
    combine :: [[SampleData]] -> [SampleData] -> [[SampleData]]
    combine int samples = map (uncurry (flip (:))) (int `zip` samples)