{-# LANGUAGE BangPatterns #-}

module ApplyStats
  ( ApplyStats(..)
  , emptyApplyStats
  , emptyApplyStatsSampled
  , applyCounted
  , runApplyCounted
  , runApplySampledWithProgress
  , runApplyGlobalCounted
  , printApplyStats
  ) where

import Research
import qualified Data.Map.Strict as M
import qualified Data.List as L
import Data.Ord (comparing)
import Data.Text (Text)
import qualified Data.Text as T
import Debug.Trace (trace)
import System.IO.Unsafe (unsafePerformIO, unsafeDupablePerformIO)
import Data.IORef

-- ---------------------------------------------------------------------------
-- Threaded stats (slow but pure)
-- ---------------------------------------------------------------------------

type Hash = Text
type AppKey = (Hash, Hash)

data ApplyStats = ApplyStats
  { totalApplyCalls :: !Int
  , uniqueApps      :: !(M.Map AppKey Int)
  , sampleInterval  :: !Int
  , sampleCounter   :: !Int
  , progressEvery   :: !Int
  }
  deriving (Show)

emptyApplyStats :: ApplyStats
emptyApplyStats = emptyApplyStatsSampled 1

emptyApplyStatsSampled :: Int -> ApplyStats
emptyApplyStatsSampled n = ApplyStats
  { totalApplyCalls = 0
  , uniqueApps      = M.empty
  , sampleInterval  = max 1 n
  , sampleCounter   = 0
  , progressEvery   = 0
  }

bump :: T -> T -> ApplyStats -> ApplyStats
bump !f !x !st =
  let !counter' = sampleCounter st + 1
      !total'   = totalApplyCalls st + 1
      !stBase   = st { totalApplyCalls = total'
                     , sampleCounter   = counter'
                     }
      !st'      = if counter' `mod` sampleInterval st /= 0
                    then stBase
                    else let !hf = termHash f
                             !hx = termHash x
                             !k  = (hf, hx)
                             !m  = M.insertWith (+) k 1 (uniqueApps st)
                         in stBase { uniqueApps = m }
  in case progressEvery st of
       0 -> st'
       n | total' `mod` n == 0 ->
            trace ("apply calls so far: " ++ show total') st'
       _ -> st'

termHash :: T -> Hash
termHash Leaf =
  nodeHash NLeaf
termHash (Stem t) =
  nodeHash (NStem (termHash t))
termHash (Fork l r) =
  nodeHash (NFork (termHash l) (termHash r))

applyCounted :: T -> T -> ApplyStats -> (T, ApplyStats)
applyCounted !f !x !st0 =
  let !st1 = bump f x st0
  in applyStepCounted f x st1

applyStepCounted :: T -> T -> ApplyStats -> (T, ApplyStats)
applyStepCounted (Fork Leaf a) _ st =
  (a, st)
applyStepCounted (Fork (Stem a) b) c st =
  let (!ac, !st1) = applyCounted a c st
      (!bc, !st2) = applyCounted b c st1
  in applyCounted ac bc st2
applyStepCounted (Fork (Fork a _b) _c) Leaf st =
  (a, st)
applyStepCounted (Fork (Fork _a b) _c) (Stem u) st =
  applyCounted b u st
applyStepCounted (Fork (Fork _a _b) c) (Fork u v) st =
  let (!cu, !st1) = applyCounted c u st
  in applyCounted cu v st1
applyStepCounted Leaf b st =
  (Stem b, st)
applyStepCounted (Stem a) b st =
  (Fork a b, st)

runApplyCounted :: T -> T -> (T, ApplyStats)
runApplyCounted !f !x =
  applyCounted f x emptyApplyStats

runApplySampled :: Int -> T -> T -> (T, ApplyStats)
runApplySampled !n !f !x =
  applyCounted f x (emptyApplyStatsSampled n)

runApplySampledWithProgress :: Int -> Int -> T -> T -> (T, ApplyStats)
runApplySampledWithProgress !interval !progress !f !x =
  let st = (emptyApplyStatsSampled interval) { progressEvery = progress }
  in applyCounted f x st

-- ---------------------------------------------------------------------------
-- Global mutable stats (fast, unsafe, single-threaded only)
-- ---------------------------------------------------------------------------

{-# NOINLINE globalTotalCount #-}
globalTotalCount :: IORef Int
globalTotalCount = unsafePerformIO (newIORef 0)

{-# NOINLINE globalInterval #-}
globalInterval :: IORef Int
globalInterval = unsafePerformIO (newIORef 1)

{-# NOINLINE globalMap #-}
globalMap :: IORef (M.Map AppKey Int)
globalMap = unsafePerformIO (newIORef M.empty)

{-# NOINLINE globalProgress #-}
globalProgress :: IORef Int
globalProgress = unsafePerformIO (newIORef 0)

resetGlobalStats :: Int -> Int -> IO ()
resetGlobalStats !interval !progress = do
  writeIORef globalTotalCount 0
  writeIORef globalInterval (max 1 interval)
  writeIORef globalMap M.empty
  writeIORef globalProgress progress

readGlobalStats :: IO ApplyStats
readGlobalStats = do
  total <- readIORef globalTotalCount
  m     <- readIORef globalMap
  pure ApplyStats
    { totalApplyCalls = total
    , uniqueApps      = m
    , sampleInterval  = 0
    , sampleCounter   = 0
    , progressEvery   = 0
    }

{-# INLINE globalBump #-}
globalBump :: T -> T -> ()
globalBump !f !x = unsafeDupablePerformIO $ do
  !total <- readIORef globalTotalCount
  let !total' = total + 1
  writeIORef globalTotalCount total'
  !interval <- readIORef globalInterval
  !progress <- readIORef globalProgress
  let !_ = if progress > 0 && total' `mod` progress == 0
            then trace ("apply calls so far: " ++ show total') ()
            else ()
  if total' `mod` interval /= 0
    then pure ()
    else do
      let !hf = termHash f
          !hx = termHash x
          !k  = (hf, hx)
      !m <- readIORef globalMap
      writeIORef globalMap (M.insertWith (+) k 1 m)
      pure ()

applyGlobalCounted :: T -> T -> T
applyGlobalCounted !f !x =
  let !_ = globalBump f x
  in applyGlobalStep f x

applyGlobalStep :: T -> T -> T
applyGlobalStep (Fork Leaf a)         _         = a
applyGlobalStep (Fork (Stem a) b) c             =
  applyGlobalCounted (applyGlobalCounted a c) (applyGlobalCounted b c)
applyGlobalStep (Fork (Fork a _b) _c)  Leaf      = a
applyGlobalStep (Fork (Fork _a b) _c) (Stem u)   = applyGlobalCounted b u
applyGlobalStep (Fork (Fork _a _b) c) (Fork u v) =
  applyGlobalCounted (applyGlobalCounted c u) v
applyGlobalStep  Leaf b                         = Stem b
applyGlobalStep (Stem a) b                      = Fork a b

runApplyGlobalCounted :: Int -> Int -> T -> T -> IO (T, ApplyStats)
runApplyGlobalCounted !interval !progress !f !x = do
  resetGlobalStats interval progress
  let !result = applyGlobalCounted f x
  !stats <- readGlobalStats
  pure (result, stats)

-- ---------------------------------------------------------------------------
-- Printing
-- ---------------------------------------------------------------------------

printApplyStats :: ApplyStats -> IO ()
printApplyStats st = do
  let !total = totalApplyCalls st
      !uniq  = M.size (uniqueApps st)
      !ratio =
        if uniq == 0
          then 0 :: Double
          else fromIntegral total / fromIntegral uniq

      counts =
        reverse
          . L.sortBy (comparing snd)
          . M.toList
          $ uniqueApps st

      repeated =
        filter ((> 1) . snd) counts

      top20 = take 20 repeated

  putStrLn $ "total apply calls: " ++ show total
  putStrLn $ "unique application patterns: " ++ show uniq
  putStrLn $ "duplication ratio total/unique: " ++ show ratio
  putStrLn $ "repeated application patterns: " ++ show (length repeated)

  putStrLn "top repeated application counts:"
  mapM_ printTop top20
  where
    short h = T.unpack (T.take 12 h)

    printTop ((hf, hx), n) =
      putStrLn $
        "  " ++ show n
        ++ "x  apply "
        ++ short hf
        ++ " "
        ++ short hx