{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}

-- The instance for NFData (TVar a) is an orphan, but necessary here
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Main where

import Control.Concurrent.STM
import Control.DeepSeq
import Control.Exception (evaluate)
import Control.Monad.Identity
import Control.Monad.IO.Class
import Criterion.Main
import Data.Dependent.Map (DMap)
import qualified Data.Dependent.Map as DMap
import Data.Dependent.Sum
import Data.Functor.Misc
import Data.IORef
import Reflex
import Reflex.Host.Class

main :: IO ()
main = defaultMain
  [ bgroup "micro" micros ]

#if !(MIN_VERSION_deepseq(1,4,2))
instance NFData (IORef a) where
  rnf x = seq x ()
#endif

instance NFData (TVar a) where
  rnf x = seq x ()

newtype WHNF a = WHNF a
instance NFData (WHNF a) where
  rnf (WHNF a) = seq a ()

withSetup :: NFData b => String -> SpiderHost Global a -> (a -> SpiderHost Global b) -> Benchmark
withSetup name setup action = env (WHNF <$> runSpiderHost setup) $ \ ~(WHNF a) ->
  bench name . nfIO $ runSpiderHost (action a)

withSetupWHNF :: String -> SpiderHost Global a -> (a -> SpiderHost Global b) -> Benchmark
withSetupWHNF name setup action = env (WHNF <$> runSpiderHost setup) $ \ ~(WHNF a) ->
  bench name . whnfIO $ runSpiderHost (action a)

micros :: [Benchmark]
micros =
  [ bench "newIORef" $ whnfIO $ void $ newIORef ()
  , env (newIORef (42 :: Int)) (bench "readIORef" . whnfIO . readIORef)
  , bench "newTVar" $ whnfIO $ void $ newTVarIO ()
  , env (newTVarIO (42 :: Int)) (bench "readTVar" . whnfIO . readTVarIO)
  , bench "newEventWithTrigger" $ whnfIO . void $ runSpiderHost $ newEventWithTrigger $
      \trigger -> return () <$ evaluate trigger
  , bench "newEventWithTriggerRef" $ whnfIO . void $ runSpiderHost newEventWithTriggerRef
  , withSetupWHNF "subscribeEvent" newEventWithTriggerRef $ subscribeEvent . fst
  , withSetupWHNF "subscribeSwitch"
    (join $ hold <$> fmap fst newEventWithTriggerRef <*> fmap fst newEventWithTriggerRef)
    (subscribeEvent . switch)
  , withSetupWHNF "subscribeMerge(1)" (setupMerge 1) $ \(ev,_) -> subscribeEvent ev
  , withSetupWHNF "subscribeMerge(100)" (setupMerge 100) (subscribeEvent . fst)
  , withSetupWHNF "subscribeMerge(10000)" (setupMerge 10000) (subscribeEvent . fst)
  , bench "runHostFrame" $ whnfIO $ runSpiderHost $ runHostFrame $ return ()
  , withSetupWHNF "fireEventsAndRead(single/single)"
    (newEventWithTriggerRef >>= subscribePair)
    (\(subd, trigger) -> fireAndRead trigger (42 :: Int) subd)
  , withSetupWHNF "fireEventsOnly"
    (newEventWithTriggerRef >>= subscribePair)
    (\(_, trigger) -> do
        Just key <- liftIO $ readIORef trigger
        fireEvents [key :=> Identity (42 :: Int)])
  , withSetupWHNF "fireEventsAndRead(head/merge1)"
    (setupMerge 1 >>= subscribePair)
    (\(subd, t:_) -> fireAndRead t (42 :: Int) subd)
  , withSetupWHNF "fireEventsAndRead(head/merge100)"
    (setupMerge 100 >>= subscribePair)
    (\(subd, t:_) -> fireAndRead t (42 :: Int) subd)
  , withSetupWHNF "fireEventsAndRead(head/merge10000)"
      (setupMerge 10000 >>= subscribePair)
      (\(subd, t:_) -> fireAndRead t (42 :: Int) subd)
  , withSetupWHNF "fireEventsOnly(head/merge100)"
    (setupMerge 100 >>= subscribePair)
    (\(_, t:_) -> do
        Just key <- liftIO $ readIORef t
        fireEvents [key :=> Identity (42 :: Int)])
  , withSetupWHNF "hold" newEventWithTriggerRef $ \(ev, _) -> hold (42 :: Int) ev
  , withSetupWHNF "sample" (newEventWithTriggerRef >>= hold (42 :: Int) . fst) sample
  ]

setupMerge :: Int
           -> SpiderHost Global ( Event (SpiderEnv Global) (DMap (Const2 Int a) Identity)
                                , [IORef (Maybe (EventTrigger Spider a))]
                                )
setupMerge num = do
  (evs, triggers) <- unzip <$> replicateM num newEventWithTriggerRef
  let !m = DMap.fromList [Const2 i :=> v | (i,v) <- zip [0..] evs]
  pure (merge m, triggers)

subscribePair :: (Event (SpiderEnv Global) a, b) -> SpiderHost Global (EventHandle (SpiderEnv Global) a, b)
subscribePair (ev, b) = (,b) <$> subscribeEvent ev

fireAndRead :: IORef (Maybe (EventTrigger (SpiderEnv Global) a)) -> a -> EventHandle (SpiderEnv Global) b
            -> SpiderHost Global (Maybe b)
fireAndRead trigger val subd = do
  Just key <- liftIO $ readIORef trigger
  fireEventsAndRead [key :=> Identity val] $ readEvent subd >>= sequence