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

{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -Wno-orphans #-}

module Main (main) where

import Data.IORef (newIORef, readIORef, modifyIORef)
import Control.DeepSeq (NFData(..))
import Data.Functor.Identity (Identity(..))
import System.Random (randomRIO)

import Streamly.Internal.Data.Stream (Stream)
import Streamly.Internal.Data.Scanl (Scanl(..))
import Streamly.Internal.Data.MutArray (MutArray)

import qualified Streamly.Internal.Data.Fold as FL
import qualified Streamly.Internal.Data.Scanl as Scanl
import qualified Streamly.Internal.Data.Stream as Stream

import Test.Tasty.Bench hiding (env)
import Streamly.Benchmark.Common
import Prelude hiding (last, length, all, any, take, unzip, sequence_, filter)

import qualified Data.Set as Set
import System.IO.Unsafe (unsafePerformIO)

-------------------------------------------------------------------------------
-- Helpers
-------------------------------------------------------------------------------

{-# INLINE source #-}
source :: (Monad m, Num a, Stream.Enumerable a) =>
    Int -> a -> Stream m a
source len from =
    Stream.enumerateFromThenTo from (from + 1) (from + fromIntegral len)

{-# INLINE benchScanWith #-}
benchScanWith :: Num a =>
    (Int -> a -> Stream IO a) -> Int -> String -> Scanl IO a b -> Benchmark
benchScanWith src len name f =
    bench name
        $ nfIO
        $ randomRIO (1, 1 :: Int)
            >>= Stream.fold FL.drain
                . Stream.postscanl f
                . src len
                . fromIntegral

{-# INLINE benchWithPostscan #-}
benchWithPostscan :: Int -> String -> Scanl IO Int a -> Benchmark
benchWithPostscan = benchScanWith source

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

moduleName :: String
moduleName = "Data.Scanl"

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

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

o_1_space_serial :: Int -> [Benchmark]
o_1_space_serial value =
    [ bgroup "key-value"
            [
              benchWithPostscan value "demuxIO (1-shot) (64 buckets) [sum 100]"
                  $ Scanl.demuxIO (getKey 64) getScanl
            , benchWithPostscan value "demuxIO (64 buckets) [sum]"
                  $ Scanl.demuxIO (getKey 64) (const (pure (Just Scanl.sum)))
            , benchWithPostscan value "classifyIO (64 buckets) [sum 100]"
                  $ Scanl.classifyIO (getKey 64) (limitedSum 100)
            , benchWithPostscan value "classifyIO (64 buckets) [sum]"
                  $ Scanl.classifyIO (getKey 64) Scanl.sum
            ]
    ]

    where

    limitedSum n = Scanl.take n Scanl.sum

    getKey buckets = (`mod` buckets)

    afterDone action (Scanl.Scanl step i e f) = Scanl.Scanl step1 i e f
        where
        step1 x a = do
            res <- step x a
            case res of
                Scanl.Partial s1 -> pure $ Scanl.Partial s1
                Scanl.Done b -> action >> pure (Scanl.Done b)

    ref = unsafePerformIO $ newIORef Set.empty
    getScanl k = do
        set <- readIORef ref
        if Set.member k set
        then pure Nothing
        else pure
                 $ Just
                 $ afterDone (modifyIORef ref (Set.insert k)) (limitedSum 100)

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

main :: IO ()
main = runWithCLIOpts defaultStreamSize allBenchmarks

    where

    allBenchmarks value =
        [ bgroup (o_1_space_prefix moduleName) (o_1_space_serial value)
        ]