-- |
-- Module      : Streamly.Benchmark.Data.Fold.Prelude
-- Copyright   : (c) 2018 Composewell
--
-- License     : MIT
-- Maintainer  : streamly@composewell.com

{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -fno-warn-warnings-deprecations #-}
{-# OPTIONS_GHC -Wno-orphans #-}

#undef FUSION_CHECK
#ifdef FUSION_CHECK
{-# OPTIONS_GHC -ddump-simpl -ddump-to-file -dsuppress-all #-}
#endif

module Main (main) where

import Control.DeepSeq (NFData(..))
import Control.Monad.IO.Class (MonadIO(..))
import Data.Functor.Identity (Identity(..))
import Data.Hashable (Hashable)
import Data.HashMap.Strict (HashMap)
import System.Random (randomRIO)

import Streamly.Internal.Data.Stream (Stream)
import Streamly.Internal.Data.Fold (Fold(..))
import Streamly.Internal.Data.IsMap.HashMap ()
import Streamly.Internal.Data.MutArray (MutArray)

import qualified Streamly.Data.Fold.Prelude as Fold
import qualified Streamly.Internal.Data.Fold as Fold
import qualified Streamly.Internal.Data.Stream as Stream

import Test.Tasty.Bench hiding (env)
import Streamly.Benchmark.Common
import Streamly.Benchmark.Common.Handle

-- We need a monadic bind here to make sure that the function f does not get
-- completely optimized out by the compiler in some cases.

{-# INLINE sourceUnfoldrM #-}
sourceUnfoldrM :: Monad m => Int -> Int -> Stream m Int
sourceUnfoldrM value n = Stream.unfoldrM step n
    where
    step cnt =
        if cnt > n + value
        then return Nothing
        else return (Just (cnt, cnt + 1))

{-# INLINE source #-}
source :: Monad m => Int -> Int -> Stream m Int
source = sourceUnfoldrM

-- | Takes a fold method, and uses it with a default source.
{-# INLINE benchIOSink #-}
benchIOSink
    :: NFData b
    => Int -> String -> (Stream IO Int -> IO b) -> Benchmark
benchIOSink value name f = bench name $ nfIO $ randomRIO (1,1) >>= f . source value

-------------------------------------------------------------------------------
-- Partitioning
-------------------------------------------------------------------------------

{-# INLINE demuxToHashMap  #-}
demuxToHashMap :: (Monad m, Hashable k
#if __GLASGOW_HASKELL__ == 810
    , Eq k
#endif
    ) =>
    (a -> k) -> (a -> m (Fold m a b)) -> Stream m a -> m (HashMap k b)
demuxToHashMap f g = Stream.fold (Fold.demuxToContainer f g)

{-# INLINE demuxToHashMapIO  #-}
demuxToHashMapIO :: (MonadIO m, Hashable k
#if __GLASGOW_HASKELL__ == 810
    , Eq k
#endif
    ) =>
    (a -> k) -> (a -> m (Fold m a b)) -> Stream m a -> m (HashMap k b)
demuxToHashMapIO f g = Stream.fold (Fold.demuxToContainerIO f g)

{-# INLINE toHashMapIO #-}
toHashMapIO :: (MonadIO m, Num a, Hashable k
#if __GLASGOW_HASKELL__ == 810
    , Eq k
#endif
    ) =>
    (a -> k) -> Stream m a -> m (HashMap k a)
toHashMapIO f = Stream.fold (Fold.toHashMapIO f Fold.sum)

-------------------------------------------------------------------------------
-- Benchmarks
-------------------------------------------------------------------------------

moduleName :: String
moduleName = "Data.Fold.Prelude"

instance NFData (MutArray a) where
    {-# INLINE rnf #-}
    rnf _ = ()

instance NFData a => NFData (Stream Identity a) where
    {-# INLINE rnf #-}
    rnf xs = runIdentity $ Stream.fold (Fold.foldl' (\_ x -> rnf x) ()) xs

o_n_heap_serial :: Int -> [Benchmark]
o_n_heap_serial value =
    [ bgroup "key-value"
            [
              benchIOSink value "demuxToHashMap (64 buckets) [sum, length]"
                $ demuxToHashMap (getKey 64) (getFold . getKey 64)
            , benchIOSink value "demuxToHashMapIO (64 buckets) [sum, length]"
                $ demuxToHashMapIO (getKey 64) (getFold . getKey 64)

            -- classify: mutable cells
            , benchIOSink value "toHashMapIO (single bucket) sum"
                $ toHashMapIO (getKey 1)
            , benchIOSink value "toHashMapIO (64 buckets) sum"
                $ toHashMapIO (getKey 64)
            , benchIOSink value "toHashMapIO (max buckets) sum"
                $ toHashMapIO (getKey value)
            ]
    ]

    where

    getKey buckets = (`mod` buckets)

    getFold k =
        return $ case k of
            0 -> Fold.sum
            1 -> Fold.length
            _ -> Fold.length

-------------------------------------------------------------------------------
-- Driver
-------------------------------------------------------------------------------

main :: IO ()
main = do
#ifndef FUSION_CHECK
    env <- mkHandleBenchEnv
    runWithCLIOpts defaultStreamSize (allBenchmarks env)

    where

    allBenchmarks _env value =
        [ bgroup (o_n_heap_prefix moduleName) (o_n_heap_serial value)
        ]
#else
    -- Enable FUSION_CHECK macro at the beginning of the file
    -- Enable one benchmark below, and run the benchmark
    -- Check the .dump-simpl output
    let value = 100000
    let input = source value 1
    let getKey buckets = (`mod` buckets)
    let getFold k =
            return $ case k of
                0 -> Fold.sum
                1 -> Fold.length
                _ -> Fold.length

    -- demuxToMap (getKey 64) (getFold . getKey 64) input
    toIntMapIO (getKey 64) input
    return ()
#endif