-- |
-- Module      : Foundation.Timing
-- License     : BSD-style
-- Maintainer  : Foundation maintainers
--
-- An implementation of a timing framework
--
{-# LANGUAGE CPP #-}
module Foundation.Timing
    ( Timing(..)
    , Measure(..)
    , stopWatch
    , measure
    ) where

import           Basement.Imports hiding (from)
import           Basement.From (from)
#if __GLASGOW_HASKELL__ < 802
import           Basement.Cast (cast)
#endif
import           Basement.Monad
-- import           Basement.UArray hiding (unsafeFreeze)
import           Basement.UArray.Mutable (MUArray)
import           Foundation.Collection
import           Foundation.Time.Types
import           Foundation.Numerical
import           Foundation.Time.Bindings
import           Control.Exception (evaluate)
import           System.Mem (performGC)
import           Data.Function (on)
import qualified GHC.Stats as GHC


data Timing = Timing
    { timeDiff           :: !NanoSeconds
    , timeBytesAllocated :: !(Maybe Word64)
    }

data Measure = Measure
    { measurements :: UArray NanoSeconds
    , iters        :: Word
    }

#if __GLASGOW_HASKELL__ >= 802
type GCStats = GHC.RTSStats

getGCStats :: IO (Maybe GCStats)
getGCStats = do
    r <- GHC.getRTSStatsEnabled
    if r then pure Nothing else Just <$> GHC.getRTSStats

diffGC :: Maybe GHC.RTSStats -> Maybe GHC.RTSStats -> Maybe Word64
diffGC gc2 gc1 = ((-) `on` GHC.allocated_bytes) <$> gc2 <*> gc1
#else
type GCStats = GHC.GCStats

getGCStats :: IO (Maybe GCStats)
getGCStats = do
    r <- GHC.getGCStatsEnabled
    if r then pure Nothing else Just <$> GHC.getGCStats

diffGC :: Maybe GHC.GCStats -> Maybe GHC.GCStats -> Maybe Word64
diffGC gc2 gc1 = cast <$> (((-) `on` GHC.bytesAllocated) <$> gc2 <*> gc1)
#endif

-- | Simple one-time measurement of time & other metrics spent in a function
stopWatch :: (a -> b) -> a -> IO Timing
stopWatch f !a = do
    performGC
    gc1 <- getGCStats
    (_, ns) <- measuringNanoSeconds (evaluate $ f a)
    gc2 <- getGCStats
    return $ Timing ns (diffGC gc2 gc1)

-- | In depth timing & other metrics analysis of a function
measure :: Word -> (a -> b) -> a -> IO Measure
measure nbIters f a = do
    d <- mutNew (from nbIters) :: IO (MUArray NanoSeconds (PrimState IO))
    loop d 0
    Measure <$> unsafeFreeze d
            <*> pure nbIters
  where
    loop d !i
        | i == nbIters = return ()
        | otherwise    = do
            (_, r) <- measuringNanoSeconds (evaluate $ f a)
            mutUnsafeWrite d (from i) r
            loop d (i+1)