{-# LANGUAGE AllowAmbiguousTypes #-}

{-|
Module      : Control.Monad.Freer.Profiling
Description : Automatic profiling of freer-simple effects
Copyright   : (c) Ben Weitzman 2018
License     : MIT
Maintainer  : ben@costarastrolgoy.com
Stability   : experimental
Portability : POSIX
-}

module Control.Monad.Freer.Profiling
  (profileEffect
  ,Named   
  ,SingleEffectStats(..)
  ,EffectStats(..)
  ,forEachStat
  ,drawEffectStats
  )
where 

import Control.Monad
import Control.Monad.Freer
import Control.Monad.Freer.Writer
import Data.Map (Map)
import qualified Data.Map as M
import Data.Tree (Forest)
import qualified Data.Tree as T
import Data.Monoid
import Data.Time.Clock.POSIX
import Text.Printf

-- | The stats relating to a single effect. Includes the number of times run and the
-- total amount of time spent running.
data SingleEffectStats = SingleEffectStats
  { effectStatCount :: Int
  , effectStatTime  :: Double
  } deriving (Show)


instance Semigroup SingleEffectStats where
    (SingleEffectStats a x) <> (SingleEffectStats b y) = SingleEffectStats (a + b) (x + y)  

instance Monoid SingleEffectStats where
    mempty = SingleEffectStats 0 0

-- | The stats for a whole bunch of effects
newtype EffectStats = EffectStats (Map String (Map String SingleEffectStats)) deriving Show

instance Semigroup EffectStats where
    (EffectStats l) <> (EffectStats r) = EffectStats $ M.unionWith (M.unionWith mappend) l r  

instance Monoid EffectStats where
    mempty = EffectStats M.empty

-- | 'Named' is a typeclass that defines the type names and constructor names for effects.
-- Effects are defined as GADTs, and unfortunately, GADTs don't work well with 'GHC.Generics'
-- or 'Data.Data', so this is a small typeclass that can be implemented manually to provide
-- the info the profiler needs.
class Named f where
    getDataTypeName :: f a -> String
    getConstructorName :: f a -> String

-- | Iterate through each indivudal stat 
forEachStat :: Monad m
            => EffectStats
            -> (String -> String -> SingleEffectStats -> m ())
            -> m ()
forEachStat (EffectStats stats) f = void $
    flip M.traverseWithKey stats $ \effect actions ->
        M.traverseWithKey (f effect) actions


toTree :: EffectStats -> Forest String
toTree (EffectStats m) = (\(a, b) -> T.Node a $ toSubtree b) <$> M.assocs m
    where
      toSubtree :: Map String SingleEffectStats -> Forest String
      toSubtree s = (\(a, b) -> T.Node (printed a b) []) <$> M.assocs s

      printed :: String -> SingleEffectStats -> String
      printed action (SingleEffectStats n t) = printf "%s: Ran %d time, taking %0.3f seconds" action n t

-- | Pretty print the effect stats
drawEffectStats :: EffectStats -> String
drawEffectStats = T.drawForest . toTree    


-- | Automatically time and count each individual effectful compoutation. An effect that
-- is interpreted to another effect will count towards both effects, flame graph style. 
profileEffect :: forall q r v
               . (Member (Writer EffectStats) r, Member IO r, Member q r, Named q)
              => Eff r v -> Eff r v
profileEffect eff = interpose handleEff eff
  where
    handleEff :: q a -> Eff r a
    handleEff action = do
      !start <- realToFrac <$> send getPOSIXTime
      v <- send action
      !end <- realToFrac <$> send getPOSIXTime
      let !diff = end - start
          effectName = getDataTypeName action
          actionName = getConstructorName action
      tell $ EffectStats (M.singleton effectName $ M.singleton actionName (SingleEffectStats 1 diff))
      return v