{-# LANGUAGE BangPatterns              #-}

{-# LANGUAGE ExistentialQuantification #-}

{-# LANGUAGE FlexibleContexts          #-}

{-# LANGUAGE MultiParamTypeClasses     #-}

{-# LANGUAGE TupleSections             #-}

{-# LANGUAGE RankNTypes                #-}

-- Collection of three benchmarks: a simple integral sum, monte carlo analysis,

-- and sliding vector.

import           Control.DeepSeq

import           Control.Monad               (foldM)

import           Control.Monad               (when, liftM)

import           Control.Monad.Codensity     (lowerCodensity)

import           Control.Monad.IO.Class      (MonadIO, liftIO)

import           Control.Monad.Trans.Class   (lift)

import           Criterion.Main

import           Data.Conduit

import           Data.Conduit.Internal       (ConduitM (..), Pipe (..))

import qualified Data.Conduit.Internal       as CI

import qualified Data.Conduit.List           as CL

import qualified Data.Foldable               as F

import           Data.Functor.Identity       (runIdentity)

import           Data.IORef

import           Data.List                   (foldl')

import           Data.Monoid                 (mempty)

import qualified Data.Sequence               as Seq

import qualified Data.Vector                 as VB

import qualified Data.Vector.Generic         as V

import qualified Data.Vector.Generic.Mutable as VM

import qualified Data.Vector.Unboxed         as VU

import           System.Environment          (withArgs)

import qualified System.Random.MWC           as MWC

import           Test.Hspec

data TestBench = TBGroup String [TestBench]

               | TBBench Benchmark

               | forall a b. (Eq b, Show b) => TBPure String a b (a -> b)

               | forall a. (Eq a, Show a) => TBIO String a (IO a)

               | forall a. (Eq a, Show a) => TBIOTest String (a -> IO ()) (IO a)

               | forall a. (Eq a, Show a) => TBIOBench String a (IO a) (IO ())

toSpec :: TestBench -> Spec

toSpec (TBGroup name tbs) = describe name $ mapM_ toSpec tbs

toSpec (TBBench _) = return ()

toSpec (TBPure name a b f) = it name $ f a `shouldBe` b

toSpec (TBIO name a f) = it name $ f >>= (`shouldBe` a)

toSpec (TBIOTest name spec f) = it name $ f >>= spec

toSpec (TBIOBench name a f _) = it name $ f >>= (`shouldBe` a)

toBench :: TestBench -> Benchmark

toBench (TBGroup name tbs) = bgroup name $ map toBench tbs

toBench (TBBench b) = b

toBench (TBPure name a _ f) = bench name $ whnf f a

toBench (TBIO name _ f) = bench name $ whnfIO f

toBench (TBIOTest name _ f) = bench name $ whnfIO f

toBench (TBIOBench name _ _ f) = bench name $ whnfIO f

runTestBench :: [TestBench] -> IO ()

runTestBench tbs = do

    withArgs [] $ hspec $ mapM_ toSpec tbs

    defaultMain $ map toBench tbs

main :: IO ()

main = runTestBench =<< sequence

    [ sumTB

    , mapSumTB

    , monteCarloTB

    , fmap (TBGroup "sliding window") $ sequence

        [ slidingWindow 10

        , slidingWindow 30

        , slidingWindow 100

        , slidingWindow 1000

        ]

    ]

-----------------------------------------------------------------------

sumTB :: IO TestBench

sumTB = do

    upperRef <- newIORef upper0

    return $ TBGroup "sum"

        [ TBPure "Data.List.foldl'" upper0 expected

            $ \upper -> foldl' (+) 0 [1..upper]

        , TBIO "Control.Monad.foldM" expected $ do

            upper <- readIORef upperRef

            foldM plusM 0 [1..upper]

        , TBPure "low level" upper0 expected $ \upper ->

            let go x !t

                    | x > upper = t

                    | otherwise = go (x + 1) (t + x)

             in go 1 0

        , TBIO "boxed vectors, I/O" expected $ do

            upper <- readIORef upperRef

            VB.foldM' plusM 0 $ VB.enumFromTo 1 upper

        , TBPure "boxed vectors" upper0 expected

            $ \upper -> VB.foldl' (+) 0 (VB.enumFromTo 1 upper)

        , TBPure "unboxed vectors" upper0 expected

            $ \upper -> VU.foldl' (+) 0 (VU.enumFromTo 1 upper)

        , TBPure "conduit, pure, fold" upper0 expected

            $ \upper -> runIdentity $ CL.enumFromTo 1 upper $$ CL.fold (+) 0

        , TBPure "conduit, pure, foldM" upper0 expected

            $ \upper -> runIdentity $ CL.enumFromTo 1 upper $$ CL.foldM plusM 0

        , TBIO "conduit, IO, fold" expected $ do

            upper <- readIORef upperRef

            CL.enumFromTo 1 upper $$ CL.fold (+) 0

        , TBIO "conduit, IO, foldM" expected $ do

            upper <- readIORef upperRef

            CL.enumFromTo 1 upper $$ CL.foldM plusM 0

        ]

  where

    upper0 = 10000 :: Int

    expected = sum [1..upper0]

    plusM x y = return $! x + y

-----------------------------------------------------------------------

mapSumTB :: IO TestBench

mapSumTB = return $ TBGroup "map + sum"

    [ TBPure "boxed vectors" upper0 expected

        $ \upper -> VB.foldl' (+) 0

                  $ VB.map (+ 1)

                  $ VB.map (* 2)

                  $ VB.enumFromTo 1 upper

    , TBPure "unboxed vectors" upper0 expected

        $ \upper -> VU.foldl' (+) 0

                  $ VU.map (+ 1)

                  $ VU.map (* 2)

                  $ VU.enumFromTo 1 upper

    , TBPure "conduit, connect1" upper0 expected $ \upper -> runIdentity

        $ CL.enumFromTo 1 upper

       $$ CL.map (* 2)

      =$= CL.map (+ 1)

      =$= CL.fold (+) 0

    , TBPure "conduit, connect2" upper0 expected $ \upper -> runIdentity

        $ CL.enumFromTo 1 upper

      =$= CL.map (* 2)

       $$ CL.map (+ 1)

      =$= CL.fold (+) 0

    , TBPure "conduit, connect3" upper0 expected $ \upper -> runIdentity

        $ CL.enumFromTo 1 upper

      =$= CL.map (* 2)

      =$= CL.map (+ 1)

       $$ CL.fold (+) 0

    , TBPure "conduit, inner fuse" upper0 expected $ \upper -> runIdentity

        $ CL.enumFromTo 1 upper

      =$= (CL.map (* 2)

      =$= CL.map (+ 1))

       $$ CL.fold (+) 0

    ]

  where

    upper0 = 10000 :: Int

    expected = sum $ map (+ 1) $ map (* 2) [1..upper0]

-----------------------------------------------------------------------

monteCarloTB :: IO TestBench

monteCarloTB = return $ TBGroup "monte carlo"

    [ TBIOTest "conduit" closeEnough $ do

        gen <- MWC.createSystemRandom

        successes <- CL.replicateM count (MWC.uniform gen)

                  $$ CL.fold (\t (x, y) ->

                                if (x*x + y*(y :: Double) < 1)

                                    then t + 1

                                    else t)

                        (0 :: Int)

        return $ fromIntegral successes / fromIntegral count * 4

    , TBIOTest "low level" closeEnough $ do

        gen <- MWC.createSystemRandom

        let go :: Int -> Int -> IO Double

            go 0 !t = return $! fromIntegral t / fromIntegral count * 4

            go i !t = do

                (x, y) <- MWC.uniform gen

                let t'

                        | x*x + y*(y :: Double) < 1 = t + 1

                        | otherwise = t

                go (i - 1) t'

        go count (0 :: Int)

    ]

  where

    count = 100000 :: Int

    closeEnough x

        | abs (x - 3.14159 :: Double) < 0.2 = return ()

        | otherwise = error $ "Monte carlo analysis too inaccurate: " ++ show x

-----------------------------------------------------------------------

slidingWindow :: Int -> IO TestBench

slidingWindow window = do

    upperRef <- newIORef upper0

    return $ TBGroup (show window)

        [ TBIOBench "low level, Seq" expected

            (swLowLevelSeq window upperRef id (\x y -> x . (F.toList y:)) ($ []))

            (swLowLevelSeq window upperRef () (\() y -> rnf y) id)

        , TBIOBench "conduit, Seq" expected

            (swConduitSeq window upperRef id (\x y -> x . (F.toList y:)) ($ []))

            (swConduitSeq window upperRef () (\() y -> rnf y) id)

        {- https://ghc.haskell.org/trac/ghc/ticket/9446

        , TBIOBench "low level, boxed Vector" expected

            (swLowLevelVector window upperRef id (\x y -> x . (VB.toList y:)) ($ []))

            (swLowLevelVector window upperRef () (\() y -> rnf (y :: VB.Vector Int)) id)

            -}

        , TBBench $ bench "low level, boxed Vector" $ whnfIO $

            swLowLevelVector window upperRef () (\() y -> rnf (y :: VB.Vector Int)) id

        {- https://ghc.haskell.org/trac/ghc/ticket/9446

        , TBIOBench "conduit, boxed Vector" expected

            (swConduitVector window upperRef id (\x y -> x . (VB.toList y:)) ($ []))

            (swConduitVector window upperRef () (\() y -> rnf (y :: VB.Vector Int)) id)

        -}

        , TBBench $ bench "conduit, boxed Vector" $ whnfIO $

            swConduitVector window upperRef () (\() y -> rnf (y :: VB.Vector Int)) id

        , TBIOBench "low level, unboxed Vector" expected

            (swLowLevelVector window upperRef id (\x y -> x . (VU.toList y:)) ($ []))

            (swLowLevelVector window upperRef () (\() y -> rnf (y :: VU.Vector Int)) id)

        , TBIOBench "conduit, unboxed Vector" expected

            (swConduitVector window upperRef id (\x y -> x . (VU.toList y:)) ($ []))

            (swConduitVector window upperRef () (\() y -> rnf (y :: VU.Vector Int)) id)

        ]

  where

    upper0 = 10000

    expected =

        loop [1..upper0]

      where

        loop input

            | length x == window = x : loop y

            | otherwise = []

          where

            x = take window input

            y = drop 1 input

swLowLevelSeq :: Int -> IORef Int -> t -> (t -> Seq.Seq Int -> t) -> (t -> t') -> IO t'

swLowLevelSeq window upperRef t0 f final = do

    upper <- readIORef upperRef

    let phase1 i !s

            | i > window = phase2 i s t0

            | otherwise = phase1 (i + 1) (s Seq.|> i)

        phase2 i !s !t

            | i > upper = t'

            | otherwise = phase2 (i + 1) s' t'

          where

            t' = f t s

            s' = Seq.drop 1 s Seq.|> i

    return $! final $! phase1 1 mempty

swLowLevelVector :: V.Vector v Int

                 => Int

                 -> IORef Int

                 -> t

                 -> (t -> v Int -> t)

                 -> (t -> t')

                 -> IO t'

swLowLevelVector window upperRef t0 f final = do

    upper <- readIORef upperRef

    let go !i !t _ _ _ | i > upper = return $! final $! t

        go !i !t !end _mv mv2 | end == bufSz  = newBuf >>= go i t sz mv2

        go !i !t !end mv mv2 = do

            VM.unsafeWrite mv end i

            when (end > sz) $ VM.unsafeWrite mv2 (end - sz) i

            let end' = end + 1

            t' <-

                if end' < sz

                    then return t

                    else do

                        v <- V.unsafeFreeze $ VM.unsafeSlice (end' - sz) sz mv

                        return $! f t v

            go (i + 1) t' end' mv mv2

    mv <- newBuf

    mv2 <- newBuf

    go 1 t0 0 mv mv2

  where

    sz = window

    bufSz = 2 * window

    newBuf = VM.new bufSz

swConduitSeq :: Int

             -> IORef Int

             -> t

             -> (t -> Seq.Seq Int -> t)

             -> (t -> t')

             -> IO t'

swConduitSeq window upperRef t0 f final = do

    upper <- readIORef upperRef

    t <- CL.enumFromTo 1 upper

        $= slidingWindowC window

        $$ CL.fold f t0

    return $! final t

swConduitVector :: V.Vector v Int

                => Int

                -> IORef Int

                -> t

                -> (t -> v Int -> t)

                -> (t -> t')

                -> IO t'

swConduitVector window upperRef t0 f final = do

    upper <- readIORef upperRef

    t <- CL.enumFromTo 1 upper

        $= slidingVectorC window

        $$ CL.fold f t0

    return $! final t

slidingWindowC :: Monad m => Int -> Conduit a m (Seq.Seq a)

slidingWindowC = slidingWindowCC

{-# INLINE [0] slidingWindowC #-}

{-# RULES "unstream slidingWindowC"

    forall i. slidingWindowC i = CI.unstream (CI.streamConduit (slidingWindowCC i) (slidingWindowS i))

  #-}

slidingWindowCC :: Monad m => Int -> Conduit a m (Seq.Seq a)

slidingWindowCC sz =

    go sz mempty

  where

    goContinue st = await >>=

                    maybe (return ())

                          (\x -> do

                             let st' = st Seq.|> x

                             yield st' >> goContinue (Seq.drop 1 st')

                          )

    go 0 st = yield st >> goContinue (Seq.drop 1 st)

    go !n st = CL.head >>= \m ->

               case m of

                 Nothing | n < sz -> yield st

                         | otherwise -> return ()

                 Just x -> go (n-1) (st Seq.|> x)

{-# INLINE slidingWindowCC #-}

slidingWindowS :: Monad m => Int -> CI.Stream m a () -> CI.Stream m (Seq.Seq a) ()

slidingWindowS sz (CI.Stream step ms0) =

    CI.Stream step' $ liftM (\s -> Left (s, sz, mempty)) ms0

  where

    step' (Left (s, 0, st)) = return $ CI.Emit (Right (s, st)) st

    step' (Left (s, i, st)) = do

        res <- step s

        return $ case res of

            CI.Stop () -> CI.Stop ()

            CI.Skip s' -> CI.Skip $ Left (s', i, st)

            CI.Emit s' a -> CI.Skip $ Left (s', i - 1, st Seq.|> a)

    step' (Right (s, st)) = do

        res <- step s

        return $ case res of

            CI.Stop () -> CI.Stop ()

            CI.Skip s' -> CI.Skip $ Right (s', st)

            CI.Emit s' a ->

                let st' = Seq.drop 1 st Seq.|> a

                 in CI.Emit (Right (s', st')) st'

{-# INLINE slidingWindowS #-}

slidingVectorC :: V.Vector v a => Int -> Conduit a IO (v a)

slidingVectorC = slidingVectorCC

{-# INLINE [0] slidingVectorC #-}

{-# RULES "unstream slidingVectorC"

    forall i. slidingVectorC i = CI.unstream (CI.streamConduit (slidingVectorCC i) (slidingVectorS i))

  #-}

slidingVectorCC :: V.Vector v a => Int -> Conduit a IO (v a)

slidingVectorCC sz = do

    mv <- newBuf

    mv2 <- newBuf

    go 0 mv mv2

  where

    bufSz = 2 * sz

    newBuf = liftIO (VM.new bufSz)

    go !end _mv mv2 | end == bufSz  = newBuf >>= go sz mv2

    go !end mv mv2 = do

      mx <- await

      case mx of

        Nothing -> when (end > 0 && end < sz) $ do

          v <- liftIO $ V.unsafeFreeze $ VM.take end mv

          yield v

        Just x -> do

          liftIO $ do

            VM.unsafeWrite mv end x

            when (end > sz) $ VM.unsafeWrite mv2 (end - sz) x

          let end' = end + 1

          when (end' >= sz) $ do

            v <- liftIO $ V.unsafeFreeze $ VM.unsafeSlice (end' - sz) sz mv

            yield v

          go end' mv mv2

slidingVectorS :: V.Vector v a => Int -> CI.Stream IO a () -> CI.Stream IO (v a) ()

slidingVectorS sz (CI.Stream step ms0) =

    CI.Stream step' ms1

  where

    bufSz = 2 * sz

    newBuf = liftIO (VM.new bufSz)

    ms1 = do

        s <- ms0

        mv <- newBuf

        mv2 <- newBuf

        return (s, 0, mv, mv2)

    step' (_, -1, _, _) = return $ CI.Stop ()

    step' (s, end, _mv, mv2) | end == bufSz = do

        mv3 <- newBuf

        return $ CI.Skip (s, sz, mv2, mv3)

    step' (s, end, mv, mv2) = do

        res <- step s

        case res of

            CI.Stop ()

                | end > 0 && end < sz -> do

                    v <- liftIO $ V.unsafeFreeze $ VM.take end mv

                    return $ CI.Emit (s, -1, mv, mv2) v

                | otherwise -> return $ CI.Stop ()

            CI.Skip s' -> return $ CI.Skip (s', end, mv, mv2)

            CI.Emit s' x -> liftIO $ do

                VM.unsafeWrite mv end x

                when (end > sz) $ VM.unsafeWrite mv2 (end - sz) x

                let end' = end + 1

                    state = (s', end', mv, mv2)

                if end' >= sz

                    then do

                        v <- V.unsafeFreeze $ VM.unsafeSlice (end' - sz) sz mv

                        return $ CI.Emit state v

                    else return $ CI.Skip state

{-# INLINE slidingVectorS #-}