{-# LANGUAGE TypeApplications #-}
module Main where

import Control.DeepSeq
import Control.Exception (evaluate)
import Control.Monad (forM_)
import Criterion
import Criterion.Main
import qualified System.Clock as Clock
import System.Environment (getArgs)
import System.Mem (performGC)

import qualified Numeric.AD as AD

import qualified Numeric.ADDual as ADD
import Numeric.ADDual.Examples


main :: IO ()
main = do
  args <- getArgs
  case args of
    ["--neural-graph"] -> mainNeuralGraph
    _ -> mainCriterion

mainCriterion :: IO ()
mainCriterion = defaultMain
  [benchNeural 100
  ,benchNeural 180  -- rather stably 2 GCs
  ,benchNeural 500
  ,benchNeural 2000
  ]
  where
    benchNeural :: Int -> Benchmark
    benchNeural n =
      env (pure (makeNeuralInput n)) $ \input ->
        bgroup ("neural-" ++ show n)
          [bench "dual" $ nf (\inp -> ADD.gradient' fneural inp 1.0) input
          ,bench "ad" $ nf (\inp -> AD.grad fneural inp) input]

mainNeuralGraph :: IO ()
mainNeuralGraph = do
  forM_ [10, 20 .. 300] $ \n -> do
    let input = makeNeuralInput n
    _ <- evaluate (force input)

    performGC
    t1 <- Clock.getTime Clock.Monotonic
    _ <- evaluate $ force (ADD.gradient' fneural input 1.0)
    t2 <- Clock.getTime Clock.Monotonic

    performGC
    t3 <- Clock.getTime Clock.Monotonic
    _ <- evaluate $ force (AD.grad fneural input)
    t4 <- Clock.getTime Clock.Monotonic

    let diff a b = fromIntegral (Clock.toNanoSecs (Clock.diffTimeSpec a b)) / 1e9 :: Double
    putStrLn $ show n ++ " " ++ show (diff t1 t2) ++ " " ++ show (diff t3 t4)