-- |
-- Module      : Streamly.Internal.Data.Fold
-- Copyright   : (c) 2019 Composewell Technologies
--               (c) 2013 Gabriel Gonzalez
-- License     : BSD3
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- See "Streamly.Data.Fold" for an overview and
-- "Streamly.Internal.Data.Fold.Types" for design notes.
--
-- IMPORTANT: keep the signatures consistent with the folds in Streamly.Prelude

module Streamly.Internal.Data.Fold
    (
    -- * Fold Type
      Step (..)
    , Fold (..)

    -- * Constructors
    -- | Which constructor to use?
    --
    -- * @foldl*@: If the fold never terminates i.e. does not use the 'Done'
    -- constructor otherwise use the @mkFold@ variants.
    -- * @*M@: Use the @M@ suffix variants if any of the step, initial, or
    -- extract function is monadic, otherwise use the pure variants.
    -- * @*_@: Use the @_@ suffix variants if the extract function is 'id' or
    -- 'return'.
    --
    , foldl'
    , foldlM'
    , foldl1'
    , foldt'
    , foldtM'
    , foldr
    , foldrM

    -- * Mappers
    -- | Monadic functions useful with mapM/lmapM on folds or streams.
    , tracing
    , trace

    -- * Folds
    -- ** Identity
    , fromPure
    , fromEffect

    -- ** Accumulators
    -- *** Semigroups and Monoids
    , sconcat
    , mconcat
    , foldMap
    , foldMapM

    -- *** Reducers
    , drain
    , drainBy
    , the
    , length
    , genericLength
    , countDistinct
    , countDistinctInt
    , frequency
    , mean
    , rollingHash
    , rollingHashWithSalt
    , rollingHashFirstN
    -- , rollingHashLastN

    -- *** Saturating Reducers
    -- | 'product' terminates if it becomes 0. Other folds can theoretically
    -- saturate on bounded types, and therefore terminate, however, they will
    -- run forever on unbounded types like Integer/Double.
    , sum
    , product
    , maximumBy
    , maximum
    , minimumBy
    , minimum

    -- *** Collectors
    -- | Avoid using these folds in scalable or performance critical
    -- applications, they buffer all the input in GC memory which can be
    -- detrimental to performance if the input is large.
    , toList
    , toListRev
    -- $toListRev
    , toStream
    , toStreamRev
    , toMap
    , topBy
    , top
    , bottom

    -- *** Scanners
    -- | Stateful transformation of the elements. Useful in combination with
    -- the 'scanMaybe' combinator. For scanners the result of the fold is
    -- usually a transformation of the current element rather than an
    -- aggregation of all elements till now.
    , last -- XXX prev
 -- , nthLast -- using Ring array
    , indexingWith
    , indexing
    , indexingRev
    , rollingMapM

    -- *** Filters
    -- | Useful in combination with the 'scanMaybe' combinator.
    , filtering
    , deleteBy
    , nub
    , nubInt
    , uniqBy
    , uniq
    , repeated
    , findIndices
    , elemIndices

    -- ** Terminating Folds
    -- Element folds. Terminate after inspecting one element. All these can be
    -- implemented in terms of the 'maybe' fold.
    , one
    , null
    , satisfy
    , maybe

    -- Sequence folds. Terminate after inspecting a sequence of elements.
    , drainN
    -- , lastN
    -- , (!!)
    , genericIndex
    , index
    , findM
    , find
    , lookup
    , findIndex
    , elemIndex
    , elem
    , notElem
    , all
    , any
    , and
    , or

    -- ** Trimmers
    -- | Useful in combination with the 'scanMaybe' combinator.
    , taking
    , dropping
    , takingEndByM
    , takingEndBy
    , takingEndByM_
    , takingEndBy_
    , droppingWhileM
    , droppingWhile
    , prune

    -- * Combinators
    -- ** Utilities
    , with

    -- ** Transforming the Monad
    , hoist
    , generally

    -- ** Mapping on output
    , rmapM

    -- ** Mapping on Input
    , transform
    , lmap
    --, lsequence
    , lmapM

    -- ** Sliding Window
    , slide2

    -- ** Scanning Input
    , scan
    , scanMany
    , postscan
    , indexed

    -- ** Zipping Input
    , zipWithM
    , zip

    -- ** Filtering Input
    , catMaybes
    , mapMaybeM
    , mapMaybe
    , scanMaybe
    , filter
    , filterM
    , sampleFromthen

    -- Either streams
    , lefts
    , rights
    , both

    {-
    -- ** Insertion
    -- | Insertion adds more elements to the stream.

    , insertBy
    , intersperseM

    -- ** Reordering
    , reverse
    -}

    -- ** Trimming
    , take

    -- By elements
    , takeEndBy
    , takeEndBy_
    , takeEndBySeq
    , takeEndBySeq_
    {-
    , drop
    , dropWhile
    , dropWhileM
    -}

    -- ** Serial Append
    , serialWith
    , serial_
    -- , tail
    -- , init
    , splitAt -- spanN
    -- , splitIn -- sessionN

    -- ** Parallel Distribution
    , teeWith
    , tee
    , teeWithFst
    , teeWithMin
    , distribute
    -- , distributeFst
    -- , distributeMin

    -- ** Unzipping
    , unzip
    -- These two can be expressed using lmap/lmapM and unzip
    , unzipWith
    , unzipWithM
    , unzipWithFstM
    , unzipWithMinM

    -- ** Parallel Alternative
    , shortest
    , longest

    -- ** Partitioning
    , partitionByM
    , partitionByFstM
    , partitionByMinM
    , partitionBy
    , partition

    -- ** Demultiplexing
    -- | Direct values in the input stream to different folds using an n-ary
    -- fold selector. 'demux' is a generalization of 'classify' (and
    -- 'partition') where each key of the classifier can use a different fold.
    , demux
    , demuxWith
    , demuxScanWith
    , demuxScanMutWith
    , demuxMutWith

    -- TODO: These can be implemented using the above operations
    -- , demuxWithSel -- Stop when the fold for the specified key stops
    -- , demuxWithMin -- Stop when any of the folds stop
    -- , demuxWithAll -- Stop when all the folds stop (run once)

    -- ** Classifying
    -- | In an input stream of key value pairs fold values for different keys
    -- in individual output buckets using the given fold. 'classify' is a
    -- special case of 'demux' where all the branches of the demultiplexer use
    -- the same fold.
    --
    -- Different types of maps can be used with these combinators via the IsMap
    -- type class. Hashmap performs better when there are more collisions, trie
    -- Map performs better otherwise. Trie has an advantage of sorting the keys
    -- at the same time.  For example if we want to store a dictionary of words
    -- and their meanings then trie Map would be better if we also want to
    -- display them in sorted order.

    , classify
    , classifyWith
    , classifyMutWith
    , classifyScanWith
    -- , classifyWithSel
    -- , classifyWithMin

    -- ** Splitting
    , many
    , chunksOf
    , chunksBetween

    -- ** Nesting
    , unfoldMany
    , concatSequence
    , concatMap

    -- * Running A Fold
    -- | Normally you would run a fold to completion by supplying it a stream,
    -- e.g. using 'Stream.fold'. However, you could also run a fold partially
    -- by using 'duplicate' on it and then running it with a stream.
    -- Alternatively, 'initialize', 'snoc' and 'finish' can be used to run a
    -- fold incrementally, however, that may not be the most efficient way to
    -- run a fold.
    , initialize
    , snoc
    , duplicate
    , finish

    -- * Deprecated
    , head
    , sequence
    , mapM
    , variance
    , stdDev
    )
where

#include "inline.hs"
#include "ArrayMacros.h"

import Control.Monad (void)
import Control.Monad.IO.Class (MonadIO(..))
import Data.Bifunctor (first)
import Data.Bits (shiftL, shiftR, (.|.), (.&.))
import Data.Either (isLeft, isRight, fromLeft, fromRight)
import Data.Functor.Identity (Identity(..))
import Data.Int (Int64)
import Data.IORef (newIORef, readIORef, writeIORef)
import Data.Map.Strict (Map)
import Data.Word (Word32)
import Foreign.Storable (peek, sizeOf)
import Streamly.Internal.Data.IsMap (IsMap(..))
import Streamly.Internal.Data.Maybe.Strict (Maybe'(..), toMaybe)
import Streamly.Internal.Data.Pipe.Type (Pipe (..), PipeState(..))
import Streamly.Internal.Data.Unboxed (Unboxed)
import Streamly.Internal.Data.Unfold.Type (Unfold(..))
import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
import Streamly.Internal.Data.Stream.Type (Stream)

import qualified Data.IntSet as IntSet
import qualified Data.Set as Set
import qualified Prelude
import qualified Streamly.Internal.Data.Array.Unboxed.Mut.Type as MA
import qualified Streamly.Internal.Data.Array.Unboxed.Type as Array
import qualified Streamly.Internal.Data.Fold.Window as FoldW
import qualified Streamly.Internal.Data.IsMap as IsMap
import qualified Streamly.Internal.Data.Pipe.Type as Pipe
import qualified Streamly.Internal.Data.Ring.Unboxed as Ring
import qualified Streamly.Internal.Data.Stream.StreamD.Type as StreamD
import qualified Streamly.Internal.Data.Stream.Type as Stream

import Prelude hiding
       ( filter, foldl1, drop, dropWhile, take, takeWhile, zipWith
       , foldl, foldr, map, mapM_, sequence, all, any, sum, product, elem
       , notElem, maximum, minimum, head, last, tail, length, null
       , reverse, iterate, init, and, or, lookup, (!!)
       , scanl, scanl1, replicate, concatMap, mconcat, foldMap, unzip
       , span, splitAt, break, mapM, zip, maybe)
import Streamly.Internal.Data.Fold.Type

-- $setup
-- >>> :m
-- >>> :set -package streamly
-- >>> import Prelude hiding (break, map, span, splitAt)
-- >>> import qualified Streamly.Data.Array.Unboxed as Array
-- >>> import qualified Streamly.Data.Fold as Fold
-- >>> import qualified Streamly.Data.Stream as Stream
-- >>> import qualified Streamly.Internal.Data.Array.Unboxed.Mut.Type as MA
-- >>> import qualified Streamly.Internal.Data.Fold as Fold
-- >>> import qualified Streamly.Internal.Data.Fold.Type as Fold
-- >>> import qualified Streamly.Internal.Data.Fold.Window as FoldW
-- >>> import qualified Streamly.Internal.Data.Parser as Parser
-- >>> import qualified Streamly.Internal.Data.Stream.Type as Stream
-- >>> import qualified Streamly.Internal.Data.Unfold as Unfold
-- >>> import Streamly.Internal.Data.Stream.Type (Stream)
-- >>> import Data.IORef (newIORef, readIORef, writeIORef)

------------------------------------------------------------------------------
-- hoist
------------------------------------------------------------------------------

-- | Change the underlying monad of a fold
--
-- /Pre-release/
hoist :: (forall x. m x -> n x) -> Fold m a b -> Fold n a b
hoist f (Fold step initial extract) =
    Fold (\x a -> f $ step x a) (f initial) (f . extract)

-- | Adapt a pure fold to any monad
--
-- > generally = Fold.hoist (return . runIdentity)
--
-- /Pre-release/
generally :: Monad m => Fold Identity a b -> Fold m a b
generally = hoist (return . runIdentity)

------------------------------------------------------------------------------
-- Transformations on fold inputs
------------------------------------------------------------------------------

-- | Flatten the monadic output of a fold to pure output.
--
-- @since 0.7.0
{-# DEPRECATED sequence "Use \"rmapM id\" instead" #-}
{-# INLINE sequence #-}
sequence :: Monad m => Fold m a (m b) -> Fold m a b
sequence = rmapM id

-- | Map a monadic function on the output of a fold.
--
-- @since 0.7.0
{-# DEPRECATED mapM "Use rmapM instead" #-}
{-# INLINE mapM #-}
mapM :: Monad m => (b -> m c) -> Fold m a b -> Fold m a c
mapM = rmapM

-- |
-- >>> mapMaybeM f = Fold.lmapM f . Fold.catMaybes
--
{-# INLINE mapMaybeM #-}
mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Fold m b r -> Fold m a r
mapMaybeM f = lmapM f . catMaybes

-- | @mapMaybe f fold@ maps a 'Maybe' returning function @f@ on the input of
-- the fold, filters out 'Nothing' elements, and return the values extracted
-- from 'Just'.
--
-- >>> mapMaybe f = Fold.lmap f . Fold.catMaybes
-- >>> mapMaybe f = Fold.mapMaybeM (return . f)
--
-- >>> f x = if even x then Just x else Nothing
-- >>> fld = Fold.mapMaybe f Fold.toList
-- >>> Stream.fold fld (Stream.enumerateFromTo 1 10)
-- [2,4,6,8,10]
--
-- @since 0.8.0
{-# INLINE mapMaybe #-}
mapMaybe :: Monad m => (a -> Maybe b) -> Fold m b r -> Fold m a r
mapMaybe f = lmap f . catMaybes

------------------------------------------------------------------------------
-- Transformations on fold inputs
------------------------------------------------------------------------------

-- | Apply a monadic function on the input and return the input.
--
-- >>> Stream.fold (Fold.lmapM (Fold.tracing print) Fold.drain) $ (Stream.enumerateFromTo (1 :: Int) 2)
-- 1
-- 2
--
-- /Pre-release/
--
{-# INLINE tracing #-}
tracing :: Monad m => (a -> m b) -> (a -> m a)
tracing f x = void (f x) >> return x

-- | Apply a monadic function to each element flowing through and discard the
-- results.
--
-- >>> Stream.fold (Fold.trace print Fold.drain) $ (Stream.enumerateFromTo (1 :: Int) 2)
-- 1
-- 2
--
-- >>> trace f = Fold.lmapM (Fold.tracing f)
--
-- /Pre-release/
{-# INLINE trace #-}
trace :: Monad m => (a -> m b) -> Fold m a r -> Fold m a r
trace f = lmapM (tracing f)

-- rename to lpipe?
--
-- | Apply a transformation on a 'Fold' using a 'Pipe'.
--
-- /Pre-release/
{-# INLINE transform #-}
transform :: Monad m => Pipe m a b -> Fold m b c -> Fold m a c
transform (Pipe pstep1 pstep2 pinitial) (Fold fstep finitial fextract) =
    Fold step initial extract

    where

    initial = first (Tuple' pinitial) <$> finitial

    step (Tuple' ps fs) x = do
        r <- pstep1 ps x
        go fs r

        where

        -- XXX use SPEC?
        go acc (Pipe.Yield b (Consume ps')) = do
            acc' <- fstep acc b
            return
                $ case acc' of
                      Partial s -> Partial $ Tuple' ps' s
                      Done b2 -> Done b2
        go acc (Pipe.Yield b (Produce ps')) = do
            acc' <- fstep acc b
            r <- pstep2 ps'
            case acc' of
                Partial s -> go s r
                Done b2 -> return $ Done b2
        go acc (Pipe.Continue (Consume ps')) =
            return $ Partial $ Tuple' ps' acc
        go acc (Pipe.Continue (Produce ps')) = do
            r <- pstep2 ps'
            go acc r

    extract (Tuple' _ fs) = fextract fs

{-# INLINE scanWith #-}
scanWith :: Monad m => Bool -> Fold m a b -> Fold m b c -> Fold m a c
scanWith isMany (Fold stepL initialL extractL) (Fold stepR initialR extractR) =
    Fold step initial extract

    where

    {-# INLINE runStep #-}
    runStep actionL sR = do
        rL <- actionL
        case rL of
            Done bL -> do
                rR <- stepR sR bL
                case rR of
                    Partial sR1 ->
                        if isMany
                        then runStep initialL sR1
                        else Done <$> extractR sR1
                    Done bR -> return $ Done bR
            Partial sL -> do
                !b <- extractL sL
                rR <- stepR sR b
                return
                    $ case rR of
                        Partial sR1 -> Partial (sL, sR1)
                        Done bR -> Done bR

    initial = do
        r <- initialR
        case r of
            Partial sR -> runStep initialL sR
            Done b -> return $ Done b

    step (sL, sR) x = runStep (stepL sL x) sR

    extract = extractR . snd

-- | Scan the input of a 'Fold' to change it in a stateful manner using another
-- 'Fold'. The scan stops as soon as the fold terminates.
--
-- /Pre-release/
{-# INLINE scan #-}
scan :: Monad m => Fold m a b -> Fold m b c -> Fold m a c
scan = scanWith False

-- XXX This does not fuse beacuse of the recursive step. Need to investigate.
--
-- | Scan the input of a 'Fold' to change it in a stateful manner using another
-- 'Fold'. The scan restarts with a fresh state if the fold terminates.
--
-- /Pre-release/
{-# INLINE scanMany #-}
scanMany :: Monad m => Fold m a b -> Fold m b c -> Fold m a c
scanMany = scanWith True

------------------------------------------------------------------------------
-- Filters
------------------------------------------------------------------------------

-- | A filtering scan that deletes the first occurrence of the element in the
-- stream that satisfies the given equality predicate.
--
-- >>> input = Stream.unfold Unfold.fromList [1,3,3,5]
-- >>> Stream.fold Fold.toList $ Stream.scanMaybe (Fold.deleteBy (==) 3) input
-- [1,3,5]
--
{-# INLINE_NORMAL deleteBy #-}
deleteBy :: Monad m => (a -> a -> Bool) -> a -> Fold m a (Maybe a)
deleteBy eq x0 = fmap extract $ foldl' step (Tuple' False Nothing)

    where

    step (Tuple' False _) x =
        if eq x x0
        then Tuple' True Nothing
        else Tuple' False (Just x)
    step (Tuple' True _) x = Tuple' True (Just x)

    extract (Tuple' _ x) = x

-- | Provide a sliding window of length 2 elements.
--
-- See "Streamly.Internal.Data.Fold.Window".
--
{-# INLINE slide2 #-}
slide2 :: Monad m => Fold m (a, Maybe a) b -> Fold m a b
slide2 (Fold step1 initial1 extract1) = Fold step initial extract

    where

    initial =
        first (Tuple' Nothing) <$> initial1

    step (Tuple' prev s) cur =
        first (Tuple' (Just cur)) <$> step1 s (cur, prev)

    extract (Tuple' _ s) = extract1 s

-- | Drop repeated elements that are adjacent to each other using the supplied
-- comparison function.
--
-- To strip duplicate path separators:
--
-- >>> input = Stream.unfold Unfold.fromList "//a//b"
-- >>> f x y = x == '/' && y == '/'
-- >>> Stream.fold Fold.toList $ Stream.scanMaybe (Fold.uniqBy f) input
-- "/a/b"
--
-- Space: @O(1)@
--
-- See also: 'nubBy'.
--
-- /Pre-release/
--
{-# INLINE uniqBy #-}
uniqBy :: Monad m => (a -> a -> Bool) -> Fold m a (Maybe a)
uniqBy eq = rollingMap f

    where

    f pre curr =
        case pre of
            Nothing -> Just curr
            Just x -> if x `eq` curr then Nothing else Just curr

-- | Drop repeated elements that are adjacent to each other.
--
-- >>> uniq = Fold.uniqBy (==)
--
{-# INLINE uniq #-}
uniq :: (Monad m, Eq a) => Fold m a (Maybe a)
uniq = uniqBy (==)

-- | Strip all leading and trailing occurrences of an element passing a
-- predicate and make all other consecutive occurrences uniq.
--
-- >> prune p = Stream.dropWhileAround p $ Stream.uniqBy (x y -> p x && p y)
--
-- @
-- > Stream.prune isSpace (Stream.fromList "  hello      world!   ")
-- "hello world!"
--
-- @
--
-- Space: @O(1)@
--
-- /Unimplemented/
{-# INLINE prune #-}
prune ::
    -- (Monad m, Eq a) =>
    (a -> Bool) -> Fold m a (Maybe a)
prune = error "Not implemented yet!"

-- | Emit only repeated elements, once.
--
-- /Unimplemented/
repeated :: -- (Monad m, Eq a) =>
    Fold m a (Maybe a)
repeated = error "Not implemented yet!"

-- | Used as a scan. Returns 'Just' for the first occurrence of an element,
-- returns 'Nothing' for any other occurrences.
--
-- >>> stream = Stream.fromList [1::Int,1,2,3,4,4,5,1,5,7]
-- >>> Stream.fold Fold.toList $ Stream.mapMaybe id $ Stream.postscan Fold.nub stream
-- [1,2,3,4,5,7]
--
-- /Pre-release/
{-# INLINE nub #-}
nub :: (Monad m, Ord a) => Fold m a (Maybe a)
nub = fmap (\(Tuple' _ x) -> x) $ foldl' step initial

    where

    initial = Tuple' Set.empty Nothing

    step (Tuple' set _) x =
        if Set.member x set
        then Tuple' set Nothing
        else Tuple' (Set.insert x set) (Just x)

-- | Like 'nub' but specialized to a stream of 'Int', for better performance.
--
-- /Pre-release/
{-# INLINE nubInt #-}
nubInt :: Monad m => Fold m Int (Maybe Int)
nubInt = fmap (\(Tuple' _ x) -> x) $ foldl' step initial

    where

    initial = Tuple' IntSet.empty Nothing

    step (Tuple' set _) x =
        if IntSet.member x set
        then Tuple' set Nothing
        else Tuple' (IntSet.insert x set) (Just x)

------------------------------------------------------------------------------
-- Left folds
------------------------------------------------------------------------------

------------------------------------------------------------------------------
-- Run Effects
------------------------------------------------------------------------------

-- |
-- > drainBy f = lmapM f drain
-- > drainBy = Fold.foldMapM (void . f)
--
-- Drain all input after passing it through a monadic function. This is the
-- dual of mapM_ on stream producers.
--
-- See also: 'Streamly.Prelude.mapM_'
--
-- @since 0.7.0
{-# INLINE drainBy #-}
drainBy ::  Monad m => (a -> m b) -> Fold m a ()
drainBy f = lmapM f drain

-- | Extract the last element of the input stream, if any.
--
-- > last = fmap getLast $ Fold.foldMap (Last . Just)
--
-- @since 0.7.0
{-# INLINE last #-}
last :: Monad m => Fold m a (Maybe a)
last = foldl1' (\_ x -> x)

-- | If the stream consists of one or more occurences of the same element
-- then returns that element else terminates and returns 'Nothing'.
--
-- /Pre-release/
{-# INLINE the #-}
the :: (Monad m, Eq a) => Fold m a (Maybe a)
the = foldt' step initial id

    where

    initial = Partial Nothing

    step Nothing x = Partial (Just x)
    step old@(Just x0) x =
            if x0 == x
            then Partial old
            else Done Nothing

------------------------------------------------------------------------------
-- To Summary
------------------------------------------------------------------------------

-- | Like 'length', except with a more general 'Num' return value
--
-- > genericLength = fmap getSum $ foldMap (Sum . const  1)
--
-- /Pre-release/
{-# INLINE genericLength #-}
genericLength :: (Monad m, Num b) => Fold m a b
genericLength = foldl' (\n _ -> n + 1) 0

-- | Determine the length of the input stream.
--
-- > length = fmap getSum $ Fold.foldMap (Sum . const  1)
--
-- @since 0.7.0
{-# INLINE length #-}
length :: Monad m => Fold m a Int
length = genericLength

-- XXX Try Hash set
-- XXX Add a countDistinct window fold
-- XXX Add a bloom filter fold

-- | Count non-duplicate elements in the stream.
--
-- Equivalent to using 'nub' followed by 'length' on a stream.
--
-- >>> countDistinct = Fold.postscan Fold.nub $ Fold.catMaybes $ Fold.length
--
-- The memory used is proportional to the number of distinct elements in the
-- stream, to guard against using too much memory use it as a scan and
-- terminate if the count reaches more than a threshold.
--
-- /Space/: \(\mathcal{O}(n)\)
--
-- /Pre-release/
--
{-# INLINE countDistinct #-}
countDistinct :: (Monad m, Ord a) => Fold m a Int
countDistinct = postscan nub $ catMaybes length
{-
countDistinct = fmap (\(Tuple' _ n) -> n) $ foldl' step initial

    where

    initial = Tuple' Set.empty 0

    step (Tuple' set n) x = do
        if Set.member x set
        then
            Tuple' set n
        else
            let cnt = n + 1
             in Tuple' (Set.insert x set) cnt
-}

-- | Like 'countDistinct' but specialized to a stream of 'Int', for better
-- performance.
--
-- >>> countDistinctInt = Fold.postscan Fold.nubInt $ Fold.catMaybes $ Fold.length
--
-- /Pre-release/
{-# INLINE countDistinctInt #-}
countDistinctInt :: Monad m => Fold m Int Int
countDistinctInt = postscan nubInt $ catMaybes length
{-
countDistinctInt = fmap (\(Tuple' _ n) -> n) $ foldl' step initial

    where

    initial = Tuple' IntSet.empty 0

    step (Tuple' set n) x = do
        if IntSet.member x set
        then
            Tuple' set n
        else
            let cnt = n + 1
             in Tuple' (IntSet.insert x set) cnt
 -}

-- | Determine the sum of all elements of a stream of numbers. Returns additive
-- identity (@0@) when the stream is empty. Note that this is not numerically
-- stable for floating point numbers.
--
-- >>> sum = FoldW.cumulative FoldW.sum
--
-- Same as following but numerically stable:
--
-- >>> sum = Fold.foldl' (+) 0
-- >>> sum = fmap Data.Monoid.getSum $ Fold.foldMap Data.Monoid.Sum
--
-- @since 0.7.0
{-# INLINE sum #-}
sum :: (Monad m, Num a) => Fold m a a
sum = FoldW.cumulative FoldW.sum

-- | Determine the product of all elements of a stream of numbers. Returns
-- multiplicative identity (@1@) when the stream is empty. The fold terminates
-- when it encounters (@0@) in its input.
--
-- Compare with @Fold.foldMap Product@.
--
-- @since 0.7.0
-- /Since 0.8.0 (Added 'Eq' constraint)/
{-# INLINE product #-}
product :: (Monad m, Num a, Eq a) => Fold m a a
product =  foldt' step (Partial 1) id

    where

    step x a =
        if a == 0
        then Done 0
        else Partial $ x * a

------------------------------------------------------------------------------
-- To Summary (Maybe)
------------------------------------------------------------------------------

-- | Determine the maximum element in a stream using the supplied comparison
-- function.
--
-- @since 0.7.0
{-# INLINE maximumBy #-}
maximumBy :: Monad m => (a -> a -> Ordering) -> Fold m a (Maybe a)
maximumBy cmp = foldl1' max'

    where

    max' x y =
        case cmp x y of
            GT -> x
            _ -> y

-- |
-- @
-- maximum = Fold.maximumBy compare
-- @
--
-- Determine the maximum element in a stream.
--
-- Compare with @Fold.foldMap Max@.
--
-- @since 0.7.0
{-# INLINE maximum #-}
maximum :: (Monad m, Ord a) => Fold m a (Maybe a)
maximum = foldl1' max

-- | Computes the minimum element with respect to the given comparison function
--
-- @since 0.7.0
{-# INLINE minimumBy #-}
minimumBy :: Monad m => (a -> a -> Ordering) -> Fold m a (Maybe a)
minimumBy cmp = foldl1' min'

    where

    min' x y =
        case cmp x y of
            GT -> y
            _ -> x

-- | Determine the minimum element in a stream using the supplied comparison
-- function.
--
-- @
-- minimum = 'minimumBy' compare
-- @
--
-- Compare with @Fold.foldMap Min@.
--
-- @since 0.7.0
{-# INLINE minimum #-}
minimum :: (Monad m, Ord a) => Fold m a (Maybe a)
minimum = foldl1' min

------------------------------------------------------------------------------
-- To Summary (Statistical)
------------------------------------------------------------------------------

-- | Compute a numerically stable arithmetic mean of all elements in the input
-- stream.
--
-- @since 0.7.0
{-# INLINE mean #-}
mean :: (Monad m, Fractional a) => Fold m a a
mean = fmap done $ foldl' step begin

    where

    begin = Tuple' 0 0

    step (Tuple' x n) y =
        let n1 = n + 1
         in Tuple' (x + (y - x) / n1) n1

    done (Tuple' x _) = x

-- | Compute a numerically stable (population) variance over all elements in
-- the input stream.
--
-- @since 0.7.0
{-# DEPRECATED variance "Use the streamly-statistics package instead" #-}
{-# INLINE variance #-}
variance :: (Monad m, Fractional a) => Fold m a a
variance = fmap done $ foldl' step begin

    where

    begin = Tuple3' 0 0 0

    step (Tuple3' n mean_ m2) x = Tuple3' n' mean' m2'

        where

        n' = n + 1
        mean' = (n * mean_ + x) / (n + 1)
        delta = x - mean_
        m2' = m2 + delta * delta * n / (n + 1)

    done (Tuple3' n _ m2) = m2 / n

-- | Compute a numerically stable (population) standard deviation over all
-- elements in the input stream.
--
-- @since 0.7.0
{-# DEPRECATED stdDev "Use the streamly-statistics package instead" #-}
{-# INLINE stdDev #-}
stdDev :: (Monad m, Floating a) => Fold m a a
stdDev = sqrt <$> variance

-- | Compute an 'Int' sized polynomial rolling hash
--
-- > H = salt * k ^ n + c1 * k ^ (n - 1) + c2 * k ^ (n - 2) + ... + cn * k ^ 0
--
-- Where @c1@, @c2@, @cn@ are the elements in the input stream and @k@ is a
-- constant.
--
-- This hash is often used in Rabin-Karp string search algorithm.
--
-- See https://en.wikipedia.org/wiki/Rolling_hash
--
-- @since 0.8.0
{-# INLINE rollingHashWithSalt #-}
rollingHashWithSalt :: (Monad m, Enum a) => Int64 -> Fold m a Int64
rollingHashWithSalt = foldl' step

    where

    k = 2891336453 :: Int64

    step cksum a = cksum * k + fromIntegral (fromEnum a)

-- | A default salt used in the implementation of 'rollingHash'.
{-# INLINE defaultSalt #-}
defaultSalt :: Int64
defaultSalt = -2578643520546668380

-- | Compute an 'Int' sized polynomial rolling hash of a stream.
--
-- > rollingHash = Fold.rollingHashWithSalt defaultSalt
--
-- @since 0.8.0
{-# INLINE rollingHash #-}
rollingHash :: (Monad m, Enum a) => Fold m a Int64
rollingHash = rollingHashWithSalt defaultSalt

-- | Compute an 'Int' sized polynomial rolling hash of the first n elements of
-- a stream.
--
-- > rollingHashFirstN = Fold.take n Fold.rollingHash
--
-- /Pre-release/
{-# INLINE rollingHashFirstN #-}
rollingHashFirstN :: (Monad m, Enum a) => Int -> Fold m a Int64
rollingHashFirstN n = take n rollingHash

-- XXX Compare this with the implementation in Fold.Window, preferrably use the
-- latter if performance is good.

-- | Apply a function on every two successive elements of a stream. The first
-- argument of the map function is the previous element and the second argument
-- is the current element. When processing the very first element in the
-- stream, the previous element is 'Nothing'.
--
-- /Pre-release/
--
{-# INLINE rollingMapM #-}
rollingMapM :: Monad m => (Maybe a -> a -> m b) -> Fold m a b
rollingMapM f = Fold step initial extract

    where

    -- XXX We need just a postscan. We do not need an initial result here.
    -- Or we can supply a default initial result as an argument to rollingMapM.
    initial = return $ Partial (Nothing, error "Empty stream")

    step (prev, _) cur = do
        x <- f prev cur
        return $ Partial (Just cur, x)

    extract = return . snd

{-# INLINE rollingMap #-}
rollingMap :: Monad m => (Maybe a -> a -> b) -> Fold m a b
rollingMap f = rollingMapM (\x y -> return $ f x y)

------------------------------------------------------------------------------
-- Monoidal left folds
------------------------------------------------------------------------------

-- | Append the elements of an input stream to a provided starting value.
--
-- >>> Stream.fold (Fold.sconcat 10) (fmap Data.Monoid.Sum $ Stream.enumerateFromTo 1 10)
-- Sum {getSum = 65}
--
-- @
-- sconcat = Fold.foldl' (<>)
-- @
--
-- @since 0.8.0
{-# INLINE sconcat #-}
sconcat :: (Monad m, Semigroup a) => a -> Fold m a a
sconcat = foldl' (<>)

-- | Fold an input stream consisting of monoidal elements using 'mappend'
-- and 'mempty'.
--
-- >>> Stream.fold Fold.mconcat (fmap Data.Monoid.Sum $ Stream.enumerateFromTo 1 10)
-- Sum {getSum = 55}
--
-- > mconcat = Fold.sconcat mempty
--
-- @since 0.7.0
{-# INLINE mconcat #-}
mconcat ::
    ( Monad m
    , Monoid a) => Fold m a a
mconcat = sconcat mempty

-- |
-- > foldMap f = Fold.lmap f Fold.mconcat
--
-- Make a fold from a pure function that folds the output of the function
-- using 'mappend' and 'mempty'.
--
-- >>> Stream.fold (Fold.foldMap Data.Monoid.Sum) $ Stream.enumerateFromTo 1 10
-- Sum {getSum = 55}
--
-- @since 0.7.0
{-# INLINE foldMap #-}
foldMap :: (Monad m, Monoid b
    ) => (a -> b) -> Fold m a b
foldMap f = lmap f mconcat

-- |
-- > foldMapM f = Fold.lmapM f Fold.mconcat
--
-- Make a fold from a monadic function that folds the output of the function
-- using 'mappend' and 'mempty'.
--
-- >>> Stream.fold (Fold.foldMapM (return . Data.Monoid.Sum)) $ Stream.enumerateFromTo 1 10
-- Sum {getSum = 55}
--
-- @since 0.7.0
{-# INLINE foldMapM #-}
foldMapM ::  (Monad m, Monoid b) => (a -> m b) -> Fold m a b
foldMapM act = foldlM' step (pure mempty)

    where

    step m a = do
        m' <- act a
        return $! mappend m m'

------------------------------------------------------------------------------
-- To Containers
------------------------------------------------------------------------------

-- $toListRev
-- This is more efficient than 'Streamly.Internal.Data.Fold.toList'. toList is
-- exactly the same as reversing the list after 'toListRev'.

-- | Buffers the input stream to a list in the reverse order of the input.
--
-- > toListRev = Fold.foldl' (flip (:)) []
--
-- /Warning!/ working on large lists accumulated as buffers in memory could be
-- very inefficient, consider using "Streamly.Array" instead.
--
-- @since 0.8.0

--  xn : ... : x2 : x1 : []
{-# INLINE toListRev #-}
toListRev :: Monad m => Fold m a [a]
toListRev = foldl' (flip (:)) []

------------------------------------------------------------------------------
-- Partial Folds
------------------------------------------------------------------------------

-- | A fold that drains the first n elements of its input, running the effects
-- and discarding the results.
--
-- > drainN n = Fold.take n Fold.drain
--
-- /Pre-release/
{-# INLINE drainN #-}
drainN :: Monad m => Int -> Fold m a ()
drainN n = take n drain

------------------------------------------------------------------------------
-- To Elements
------------------------------------------------------------------------------

-- | Like 'index', except with a more general 'Integral' argument
--
-- /Pre-release/
{-# INLINE genericIndex #-}
genericIndex :: (Integral i, Monad m) => i -> Fold m a (Maybe a)
genericIndex i = foldt' step (Partial 0) (const Nothing)

    where

    step j a =
        if i == j
        then Done $ Just a
        else Partial (j + 1)

-- | Lookup the element at the given index.
--
-- See also: 'Streamly.Prelude.!!'
--
-- @since 0.7.0
{-# INLINE index #-}
index :: Monad m => Int -> Fold m a (Maybe a)
index = genericIndex

-- | Map a 'Maybe' returning function on the next element in the stream.
--
-- /Pre-release/
--
{-# INLINE maybe #-}
maybe :: Monad m => (a -> Maybe b) -> Fold m a (Maybe b)
maybe f = foldt' (const (Done . f)) (Partial Nothing) id

-- | Test if the next element satisfies the supplied predicate.
--
-- /Pre-release/
{-# INLINE satisfy #-}
satisfy :: Monad m => (a -> Bool) -> Fold m a (Maybe a)
satisfy f = Fold step (return $ Partial ()) (const (return Nothing))

    where

    step () a = return $ Done $ if f a then Just a else Nothing

-- Naming notes:
--
-- "head" and "next" are two alternative names for the same API. head sounds
-- apt in the context of lists but next sounds more apt in the context of
-- streams where we think in terms of generating and consuming the next element
-- rather than taking the head of some static/persistent structure.
--
-- We also want to keep the nomenclature consistent across folds and parsers,
-- "head" becomes even more unintuitive for parsers because there are two
-- possible variants viz. peek and next.
--
-- Also, the "head" fold creates confusion in situations like
-- https://github.com/composewell/streamly/issues/1404 where intuitive
-- expectation from head is to consume the entire stream and just give us the
-- head. There we want to convey the notion that we consume one element from
-- the stream and stop. The name "one" already being used in parsers for this
-- purpose sounds more apt from this perspective.
--
-- The source of confusion is perhaps due to the fact that some folds consume
-- the entire stream and others terminate early. It may have been clearer if we
-- had separate abstractions for the two use cases.

-- XXX We can possibly use "head" for the purposes of reducing the entire
-- stream to the head element i.e. take the head and drain the rest.

-- | Take one element from the stream and stop.
--
-- /Pre-release/
--
{-# INLINE one #-}
one :: Monad m => Fold m a (Maybe a)
one = maybe Just

-- | Extract the first element of the stream, if any.
--
-- >>> head = Fold.one
--
-- @since 0.7.0
{-# DEPRECATED head "Please use \"one\" instead" #-}
{-# INLINE head #-}
head :: Monad m => Fold m a (Maybe a)
head = one

-- | Returns the first element that satisfies the given predicate.
--
-- /Pre-release/
{-# INLINE findM #-}
findM :: Monad m => (a -> m Bool) -> Fold m a (Maybe a)
findM predicate = Fold step (return $ Partial ()) (const $ return Nothing)

    where

    step () a =
        let f r =
                if r
                then Done (Just a)
                else Partial ()
         in f <$> predicate a

-- | Returns the first element that satisfies the given predicate.
--
-- @since 0.7.0
{-# INLINE find #-}
find :: Monad m => (a -> Bool) -> Fold m a (Maybe a)
find p = findM (return . p)

-- | In a stream of (key-value) pairs @(a, b)@, return the value @b@ of the
-- first pair where the key equals the given value @a@.
--
-- > lookup = snd <$> Fold.find ((==) . fst)
--
-- @since 0.7.0
{-# INLINE lookup #-}
lookup :: (Eq a, Monad m) => a -> Fold m (a,b) (Maybe b)
lookup a0 = foldt' step (Partial ()) (const Nothing)

    where

    step () (a, b) =
        if a == a0
        then Done $ Just b
        else Partial ()

-- | Returns the first index that satisfies the given predicate.
--
-- @since 0.7.0
{-# INLINE findIndex #-}
findIndex :: Monad m => (a -> Bool) -> Fold m a (Maybe Int)
findIndex predicate = foldt' step (Partial 0) (const Nothing)

    where

    step i a =
        if predicate a
        then Done $ Just i
        else Partial (i + 1)

-- XXX If we have a Continue contructor we won't need a Maybe type for the
-- scan. The fold will have the ability to skip producing output.
--
-- | Returns all indices that satisfy the given predicate.
--
{-# INLINE findIndices #-}
findIndices :: Monad m => (a -> Bool) -> Fold m a (Maybe Int)
findIndices predicate =
    -- XXX implement by combining indexing and filtering scans
    fmap (either (const Nothing) Just) $ foldl' step (Left (-1))

    where

    step i a =
        if predicate a
        then Right (either id id i + 1)
        else Left (either id id i + 1)

-- | Find all the indices where the value of the element in the stream is equal
-- to the given value.
--
-- >>> elemIndices a = Fold.findIndices (== a)
--
{-# INLINE elemIndices #-}
elemIndices :: (Monad m, Eq a) => a -> Fold m a (Maybe Int)
elemIndices a = findIndices (== a)

-- | Returns the first index where a given value is found in the stream.
--
-- > elemIndex a = Fold.findIndex (== a)
--
-- @since 0.7.0
{-# INLINE elemIndex #-}
elemIndex :: (Eq a, Monad m) => a -> Fold m a (Maybe Int)
elemIndex a = findIndex (a ==)

------------------------------------------------------------------------------
-- To Boolean
------------------------------------------------------------------------------

-- | Return 'True' if the input stream is empty.
--
-- > null = fmap isJust Fold.one
--
-- @since 0.7.0
{-# INLINE null #-}
null :: Monad m => Fold m a Bool
null = foldt' (\() _ -> Done False) (Partial ()) (const True)

-- | Returns 'True' if any of the elements of a stream satisfies a predicate.
--
-- >>> Stream.fold (Fold.any (== 0)) $ Stream.fromList [1,0,1]
-- True
--
-- > any p = Fold.lmap p Fold.or
--
-- @since 0.7.0
{-# INLINE any #-}
any :: Monad m => (a -> Bool) -> Fold m a Bool
any predicate = foldt' step initial id

    where

    initial = Partial False

    step _ a =
        if predicate a
        then Done True
        else Partial False

-- | Return 'True' if the given element is present in the stream.
--
-- > elem a = Fold.any (== a)
--
-- @since 0.7.0
{-# INLINE elem #-}
elem :: (Eq a, Monad m) => a -> Fold m a Bool
elem a = any (a ==)

-- | Returns 'True' if all elements of a stream satisfy a predicate.
--
-- >>> Stream.fold (Fold.all (== 0)) $ Stream.fromList [1,0,1]
-- False
--
-- > all p = Fold.lmap p Fold.and
--
-- @since 0.7.0
{-# INLINE all #-}
all :: Monad m => (a -> Bool) -> Fold m a Bool
all predicate = foldt' step initial id

    where

    initial = Partial True

    step _ a =
        if predicate a
        then Partial True
        else Done False

-- | Returns 'True' if the given element is not present in the stream.
--
-- > notElem a = Fold.all (/= a)
--
-- @since 0.7.0
{-# INLINE notElem #-}
notElem :: (Eq a, Monad m) => a -> Fold m a Bool
notElem a = all (a /=)

-- | Returns 'True' if all elements are 'True', 'False' otherwise
--
-- > and = Fold.all (== True)
--
-- @since 0.7.0
{-# INLINE and #-}
and :: Monad m => Fold m Bool Bool
and = all (== True)

-- | Returns 'True' if any element is 'True', 'False' otherwise
--
-- > or = Fold.any (== True)
--
-- @since 0.7.0
{-# INLINE or #-}
or :: Monad m => Fold m Bool Bool
or = any (== True)

------------------------------------------------------------------------------
-- Grouping/Splitting
------------------------------------------------------------------------------

------------------------------------------------------------------------------
-- Grouping without looking at elements
------------------------------------------------------------------------------

------------------------------------------------------------------------------
-- Binary APIs
------------------------------------------------------------------------------

-- | @splitAt n f1 f2@ composes folds @f1@ and @f2@ such that first @n@
-- elements of its input are consumed by fold @f1@ and the rest of the stream
-- is consumed by fold @f2@.
--
-- >>> let splitAt_ n xs = Stream.fold (Fold.splitAt n Fold.toList Fold.toList) $ Stream.fromList xs
--
-- >>> splitAt_ 6 "Hello World!"
-- ("Hello ","World!")
--
-- >>> splitAt_ (-1) [1,2,3]
-- ([],[1,2,3])
--
-- >>> splitAt_ 0 [1,2,3]
-- ([],[1,2,3])
--
-- >>> splitAt_ 1 [1,2,3]
-- ([1],[2,3])
--
-- >>> splitAt_ 3 [1,2,3]
-- ([1,2,3],[])
--
-- >>> splitAt_ 4 [1,2,3]
-- ([1,2,3],[])
--
-- > splitAt n f1 f2 = Fold.serialWith (,) (Fold.take n f1) f2
--
-- /Internal/

{-# INLINE splitAt #-}
splitAt
    :: Monad m
    => Int
    -> Fold m a b
    -> Fold m a c
    -> Fold m a (b, c)
splitAt n fld = serialWith (,) (take n fld)

------------------------------------------------------------------------------
-- Element Aware APIs
------------------------------------------------------------------------------
--
------------------------------------------------------------------------------
-- Binary APIs
------------------------------------------------------------------------------

{-# INLINE takingEndByM #-}
takingEndByM :: Monad m => (a -> m Bool) -> Fold m a (Maybe a)
takingEndByM p = Fold step initial (return . toMaybe)

    where

    initial = return $ Partial Nothing'

    step _ a = do
        r <- p a
        return
            $ if r
              then Done $ Just a
              else Partial $ Just' a

-- |
--
-- >>> takingEndBy p = Fold.takingEndByM (return . p)
--
{-# INLINE takingEndBy #-}
takingEndBy :: Monad m => (a -> Bool) -> Fold m a (Maybe a)
takingEndBy p = takingEndByM (return . p)

{-# INLINE takingEndByM_ #-}
takingEndByM_ :: Monad m => (a -> m Bool) -> Fold m a (Maybe a)
takingEndByM_ p = Fold step initial (return . toMaybe)

    where

    initial = return $ Partial Nothing'

    step _ a = do
        r <- p a
        return
            $ if r
              then Done Nothing
              else Partial $ Just' a

-- |
--
-- >>> takingEndBy_ p = Fold.takingEndByM_ (return . p)
--
{-# INLINE takingEndBy_ #-}
takingEndBy_ :: Monad m => (a -> Bool) -> Fold m a (Maybe a)
takingEndBy_ p = takingEndByM_ (return . p)

{-# INLINE droppingWhileM #-}
droppingWhileM :: Monad m => (a -> m Bool) -> Fold m a (Maybe a)
droppingWhileM p = Fold step initial (return . toMaybe)

    where

    initial = return $ Partial Nothing'

    step Nothing' a = do
        r <- p a
        return
            $ Partial
            $ if r
              then Nothing'
              else Just' a
    step _ a = return $ Partial $ Just' a

{-# INLINE droppingWhile #-}
droppingWhile :: Monad m => (a -> Bool) -> Fold m a (Maybe a)
droppingWhile p = droppingWhileM (return . p)

-- Note: Keep this consistent with S.splitOn. In fact we should eliminate
-- S.splitOn in favor of the fold.
--
-- XXX Use Fold.many instead once it is fixed.
--
-- | Like 'takeEndBy' but drops the element on which the predicate succeeds.
--
-- >>> Stream.fold (Fold.takeEndBy_ (== '\n') Fold.toList) $ Stream.fromList "hello\nthere\n"
-- "hello"
--
-- >>> Stream.fold Fold.toList $ Stream.foldMany (Fold.takeEndBy_ (== '\n') Fold.toList) $ Stream.fromList "hello\nthere\n"
-- ["hello","there"]
--
-- > Stream.splitOnSuffix p f = Stream.foldMany (Fold.takeEndBy_ p f)
--
-- See 'Streamly.Prelude.splitOnSuffix' for more details on splitting a
-- stream using 'takeEndBy_'.
--
-- @since 0.8.0
{-# INLINE takeEndBy_ #-}
takeEndBy_ :: Monad m => (a -> Bool) -> Fold m a b -> Fold m a b
-- takeEndBy_ predicate = scanMaybe (takingEndBy_ predicate)
takeEndBy_ predicate (Fold fstep finitial fextract) =
    Fold step finitial fextract

    where

    step s a =
        if not (predicate a)
        then fstep s a
        else Done <$> fextract s

-- | Take the input, stop when the predicate succeeds taking the succeeding
-- element as well.
--
-- >>> Stream.fold (Fold.takeEndBy (== '\n') Fold.toList) $ Stream.fromList "hello\nthere\n"
-- "hello\n"
--
-- >>> Stream.fold Fold.toList $ Stream.foldMany (Fold.takeEndBy (== '\n') Fold.toList) $ Stream.fromList "hello\nthere\n"
-- ["hello\n","there\n"]
--
-- > Stream.splitWithSuffix p f = Stream.foldMany (Fold.takeEndBy p f)
--
-- See 'Streamly.Prelude.splitWithSuffix' for more details on splitting a
-- stream using 'takeEndBy'.
--
-- @since 0.8.0
{-# INLINE takeEndBy #-}
takeEndBy :: Monad m => (a -> Bool) -> Fold m a b -> Fold m a b
-- takeEndBy predicate = scanMaybe (takingEndBy predicate)
takeEndBy predicate (Fold fstep finitial fextract) =
    Fold step finitial fextract

    where

    step s a = do
        res <- fstep s a
        if not (predicate a)
        then return res
        else do
            case res of
                Partial s1 -> Done <$> fextract s1
                Done b -> return $ Done b

------------------------------------------------------------------------------
-- Binary splitting on a separator
------------------------------------------------------------------------------

data SplitOnSeqState acc a rb rh w ck =
      SplitOnSeqEmpty !acc
    | SplitOnSeqSingle !acc !a
    | SplitOnSeqWord !acc !Int !w
    | SplitOnSeqWordLoop !acc !w
    | SplitOnSeqKR !acc !Int !rb !rh
    | SplitOnSeqKRLoop !acc !ck !rb !rh

-- XXX Need to add tests for takeEndBySeq, we have tests for takeEndBySeq_ .

-- | Continue taking the input until the input sequence matches the supplied
-- sequence, taking the supplied sequence as well. If the pattern is empty this
-- acts as an identity fold.
--
-- >>> s = Stream.fromList "hello there. How are you?"
-- >>> f = Fold.takeEndBySeq (Array.fromList "re") Fold.toList
-- >>> Stream.fold f s
-- "hello there"
--
-- >>> Stream.fold Fold.toList $ Stream.foldMany f s
-- ["hello there",". How are"," you?"]
--
-- /Pre-release/
{-# INLINE takeEndBySeq #-}
takeEndBySeq :: forall m a b. (MonadIO m, Unboxed a, Enum a, Eq a) =>
       Array.Array a
    -> Fold m a b
    -> Fold m a b
takeEndBySeq patArr (Fold fstep finitial fextract) =
    Fold step initial extract

    where

    patLen = Array.length patArr

    initial = do
        res <- finitial
        case res of
            Partial acc
                | patLen == 0 ->
                    -- XXX Should we match nothing or everything on empty
                    -- pattern?
                    -- Done <$> fextract acc
                    return $ Partial $ SplitOnSeqEmpty acc
                | patLen == 1 -> do
                    pat <- liftIO $ Array.unsafeIndexIO 0 patArr
                    return $ Partial $ SplitOnSeqSingle acc pat
                | SIZE_OF(a) * patLen <= sizeOf (undefined :: Word) ->
                    return $ Partial $ SplitOnSeqWord acc 0 0
                | otherwise -> do
                    (rb, rhead) <- liftIO $ Ring.new patLen
                    return $ Partial $ SplitOnSeqKR acc 0 rb rhead
            Done b -> return $ Done b

    -- Word pattern related
    maxIndex = patLen - 1

    elemBits = SIZE_OF(a) * 8

    wordMask :: Word
    wordMask = (1 `shiftL` (elemBits * patLen)) - 1

    wordPat :: Word
    wordPat = wordMask .&. Array.foldl' addToWord 0 patArr

    addToWord wd a = (wd `shiftL` elemBits) .|. fromIntegral (fromEnum a)

    -- For Rabin-Karp search
    k = 2891336453 :: Word32
    coeff = k ^ patLen

    addCksum cksum a = cksum * k + fromIntegral (fromEnum a)

    deltaCksum cksum old new =
        addCksum cksum new - coeff * fromIntegral (fromEnum old)

    -- XXX shall we use a random starting hash or 1 instead of 0?
    -- XXX Need to keep this cached across fold calls in foldmany
    -- XXX We may need refold to inject the cached state instead of
    -- initializing the state every time.
    -- XXX Allocation of ring buffer should also be done once
    patHash = Array.foldl' addCksum 0 patArr

    step (SplitOnSeqEmpty s) x = do
        res <- fstep s x
        case res of
            Partial s1 -> return $ Partial $ SplitOnSeqEmpty s1
            Done b -> return $ Done b
    step (SplitOnSeqSingle s pat) x = do
        res <- fstep s x
        case res of
            Partial s1
                | pat /= x -> return $ Partial $ SplitOnSeqSingle s1 pat
                | otherwise -> Done <$> fextract s1
            Done b -> return $ Done b
    step (SplitOnSeqWord s idx wrd) x = do
        res <- fstep s x
        let wrd1 = addToWord wrd x
        case res of
            Partial s1
                | idx == maxIndex -> do
                    if wrd1 .&. wordMask == wordPat
                    then Done <$> fextract s1
                    else return $ Partial $ SplitOnSeqWordLoop s1 wrd1
                | otherwise ->
                    return $ Partial $ SplitOnSeqWord s1 (idx + 1) wrd1
            Done b -> return $ Done b
    step (SplitOnSeqWordLoop s wrd) x = do
        res <- fstep s x
        let wrd1 = addToWord wrd x
        case res of
            Partial s1
                | wrd1 .&. wordMask == wordPat ->
                    Done <$> fextract s1
                | otherwise ->
                    return $ Partial $ SplitOnSeqWordLoop s1 wrd1
            Done b -> return $ Done b
    step (SplitOnSeqKR s idx rb rh) x = do
        res <- fstep s x
        case res of
            Partial s1 -> do
                rh1 <- liftIO $ Ring.unsafeInsert rb rh x
                if idx == maxIndex
                then do
                    let fld = Ring.unsafeFoldRing (Ring.ringBound rb)
                    let !ringHash = fld addCksum 0 rb
                    if ringHash == patHash && Ring.unsafeEqArray rb rh1 patArr
                    then Done <$> fextract s1
                    else return $ Partial $ SplitOnSeqKRLoop s1 ringHash rb rh1
                else
                    return $ Partial $ SplitOnSeqKR s1 (idx + 1) rb rh1
            Done b -> return $ Done b
    step (SplitOnSeqKRLoop s cksum rb rh) x = do
        res <- fstep s x
        case res of
            Partial s1 -> do
                old <- liftIO $ peek rh
                rh1 <- liftIO $ Ring.unsafeInsert rb rh x
                let ringHash = deltaCksum cksum old x
                if ringHash == patHash && Ring.unsafeEqArray rb rh1 patArr
                then Done <$> fextract s1
                else return $ Partial $ SplitOnSeqKRLoop s1 ringHash rb rh1
            Done b -> return $ Done b

    extract state =
        let st =
                case state of
                    SplitOnSeqEmpty s -> s
                    SplitOnSeqSingle s _ -> s
                    SplitOnSeqWord s _ _ -> s
                    SplitOnSeqWordLoop s _ -> s
                    SplitOnSeqKR s _ _ _ -> s
                    SplitOnSeqKRLoop s _ _ _ -> s
         in fextract st

-- | Like 'takeEndBySeq' but discards the matched sequence.
--
-- /Pre-release/
--
{-# INLINE takeEndBySeq_ #-}
takeEndBySeq_ :: forall m a b. (MonadIO m, Unboxed a, Enum a, Eq a) =>
       Array.Array a
    -> Fold m a b
    -> Fold m a b
takeEndBySeq_ patArr (Fold fstep finitial fextract) =
    Fold step initial extract

    where

    patLen = Array.length patArr

    initial = do
        res <- finitial
        case res of
            Partial acc
                | patLen == 0 ->
                    -- XXX Should we match nothing or everything on empty
                    -- pattern?
                    -- Done <$> fextract acc
                    return $ Partial $ SplitOnSeqEmpty acc
                | patLen == 1 -> do
                    pat <- liftIO $ Array.unsafeIndexIO 0 patArr
                    return $ Partial $ SplitOnSeqSingle acc pat
                -- XXX Need to add tests for this case
                | SIZE_OF(a) * patLen <= sizeOf (undefined :: Word) ->
                    return $ Partial $ SplitOnSeqWord acc 0 0
                | otherwise -> do
                    (rb, rhead) <- liftIO $ Ring.new patLen
                    return $ Partial $ SplitOnSeqKR acc 0 rb rhead
            Done b -> return $ Done b

    -- Word pattern related
    maxIndex = patLen - 1

    elemBits = SIZE_OF(a) * 8

    wordMask :: Word
    wordMask = (1 `shiftL` (elemBits * patLen)) - 1

    elemMask :: Word
    elemMask = (1 `shiftL` elemBits) - 1

    wordPat :: Word
    wordPat = wordMask .&. Array.foldl' addToWord 0 patArr

    addToWord wd a = (wd `shiftL` elemBits) .|. fromIntegral (fromEnum a)

    -- For Rabin-Karp search
    k = 2891336453 :: Word32
    coeff = k ^ patLen

    addCksum cksum a = cksum * k + fromIntegral (fromEnum a)

    deltaCksum cksum old new =
        addCksum cksum new - coeff * fromIntegral (fromEnum old)

    -- XXX shall we use a random starting hash or 1 instead of 0?
    -- XXX Need to keep this cached across fold calls in foldMany
    -- XXX We may need refold to inject the cached state instead of
    -- initializing the state every time.
    -- XXX Allocation of ring buffer should also be done once
    patHash = Array.foldl' addCksum 0 patArr

    step (SplitOnSeqEmpty s) x = do
        res <- fstep s x
        case res of
            Partial s1 -> return $ Partial $ SplitOnSeqEmpty s1
            Done b -> return $ Done b
    step (SplitOnSeqSingle s pat) x = do
        if pat /= x
        then do
            res <- fstep s x
            case res of
                Partial s1 -> return $ Partial $ SplitOnSeqSingle s1 pat
                Done b -> return $ Done b
        else Done <$> fextract s
    step (SplitOnSeqWord s idx wrd) x = do
        let wrd1 = addToWord wrd x
        if idx == maxIndex
        then do
            if wrd1 .&. wordMask == wordPat
            then Done <$> fextract s
            else return $ Partial $ SplitOnSeqWordLoop s wrd1
        else return $ Partial $ SplitOnSeqWord s (idx + 1) wrd1
    step (SplitOnSeqWordLoop s wrd) x = do
        let wrd1 = addToWord wrd x
            old = (wordMask .&. wrd)
                    `shiftR` (elemBits * (patLen - 1))
        res <- fstep s (toEnum $ fromIntegral old)
        case res of
            Partial s1
                | wrd1 .&. wordMask == wordPat ->
                    Done <$> fextract s1
                | otherwise ->
                    return $ Partial $ SplitOnSeqWordLoop s1 wrd1
            Done b -> return $ Done b
    step (SplitOnSeqKR s idx rb rh) x = do
        rh1 <- liftIO $ Ring.unsafeInsert rb rh x
        if idx == maxIndex
        then do
            let fld = Ring.unsafeFoldRing (Ring.ringBound rb)
            let !ringHash = fld addCksum 0 rb
            if ringHash == patHash && Ring.unsafeEqArray rb rh1 patArr
            then Done <$> fextract s
            else return $ Partial $ SplitOnSeqKRLoop s ringHash rb rh1
        else return $ Partial $ SplitOnSeqKR s (idx + 1) rb rh1
    step (SplitOnSeqKRLoop s cksum rb rh) x = do
        old <- liftIO $ peek rh
        res <- fstep s old
        case res of
            Partial s1 -> do
                rh1 <- liftIO $ Ring.unsafeInsert rb rh x
                let ringHash = deltaCksum cksum old x
                if ringHash == patHash && Ring.unsafeEqArray rb rh1 patArr
                then Done <$> fextract s1
                else return $ Partial $ SplitOnSeqKRLoop s1 ringHash rb rh1
            Done b -> return $ Done b

    -- XXX extract should return backtrack count as well. If the fold
    -- terminates early inside extract, we may still have buffered data
    -- remaining which will be lost if we do not communicate that to the
    -- driver.
    extract state = do
        let consumeWord s n wrd = do
                if n == 0
                then fextract s
                else do
                    let old = elemMask .&. (wrd `shiftR` (elemBits * (n - 1)))
                    r <- fstep s (toEnum $ fromIntegral old)
                    case r of
                        Partial s1 -> consumeWord s1 (n - 1) wrd
                        Done b -> return b

        let consumeRing s n rb rh =
                if n == 0
                then fextract s
                else do
                    old <- liftIO $ peek rh
                    let rh1 = Ring.advance rb rh
                    r <- fstep s old
                    case r of
                        Partial s1 -> consumeRing s1 (n - 1) rb rh1
                        Done b -> return b

        case state of
            SplitOnSeqEmpty s -> fextract s
            SplitOnSeqSingle s _ -> fextract s
            SplitOnSeqWord s idx wrd -> consumeWord s idx wrd
            SplitOnSeqWordLoop s wrd -> consumeWord s patLen wrd
            SplitOnSeqKR s idx rb _ -> consumeRing s idx rb (Ring.startOf rb)
            SplitOnSeqKRLoop s _ rb rh -> consumeRing s patLen rb rh

------------------------------------------------------------------------------
-- Distributing
------------------------------------------------------------------------------
--
-- | Distribute one copy of the stream to each fold and zip the results.
--
-- @
--                 |-------Fold m a b--------|
-- ---stream m a---|                         |---m (b,c)
--                 |-------Fold m a c--------|
-- @
-- >>> Stream.fold (Fold.tee Fold.sum Fold.length) (Stream.enumerateFromTo 1.0 100.0)
-- (5050.0,100)
--
-- > tee = teeWith (,)
--
-- @since 0.7.0
{-# INLINE tee #-}
tee :: Monad m => Fold m a b -> Fold m a c -> Fold m a (b,c)
tee = teeWith (,)

-- XXX use "List" instead of "[]"?, use Array for output to scale it to a large
-- number of consumers? For polymorphic case a vector could be helpful. For
-- Storables we can use arrays. Will need separate APIs for those.
--
-- | Distribute one copy of the stream to each fold and collect the results in
-- a container.
--
-- @
--
--                 |-------Fold m a b--------|
-- ---stream m a---|                         |---m [b]
--                 |-------Fold m a b--------|
--                 |                         |
--                            ...
-- @
--
-- >>> Stream.fold (Fold.distribute [Fold.sum, Fold.length]) (Stream.enumerateFromTo 1 5)
-- [15,5]
--
-- >>> distribute = Prelude.foldr (Fold.teeWith (:)) (Fold.fromPure [])
--
-- This is the consumer side dual of the producer side 'sequence' operation.
--
-- Stops when all the folds stop.
--
-- @since 0.7.0
{-# INLINE distribute #-}
distribute :: Monad m => [Fold m a b] -> Fold m a [b]
distribute = Prelude.foldr (teeWith (:)) (fromPure [])

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

{-# INLINE partitionByMUsing #-}
partitionByMUsing :: Monad m =>
       (  (x -> y -> (x, y))
       -> Fold m (Either b c) x
       -> Fold m (Either b c) y
       -> Fold m (Either b c) (x, y)
       )
    -> (a -> m (Either b c))
    -> Fold m b x
    -> Fold m c y
    -> Fold m a (x, y)
partitionByMUsing t f fld1 fld2 =
    let l = lmap (fromLeft undefined) fld1  -- :: Fold m (Either b c) x
        r = lmap (fromRight undefined) fld2 -- :: Fold m (Either b c) y
     in lmapM f (t (,) (filter isLeft l) (filter isRight r))

-- | Partition the input over two folds using an 'Either' partitioning
-- predicate.
--
-- @
--
--                                     |-------Fold b x--------|
-- -----stream m a --> (Either b c)----|                       |----(x,y)
--                                     |-------Fold c y--------|
-- @
--
-- Send input to either fold randomly:
--
-- @
-- > import System.Random (randomIO)
-- > randomly a = randomIO >>= \\x -> return $ if x then Left a else Right a
-- > Stream.fold (Fold.partitionByM randomly Fold.length Fold.length) (Stream.enumerateFromTo 1 100)
-- (59,41)
-- @
--
-- Send input to the two folds in a proportion of 2:1:
--
-- >>> :{
-- proportionately m n = do
--  ref <- newIORef $ cycle $ concat [replicate m Left, replicate n Right]
--  return $ \a -> do
--      r <- readIORef ref
--      writeIORef ref $ tail r
--      return $ Prelude.head r a
-- :}
--
-- >>> :{
-- main = do
--  f <- proportionately 2 1
--  r <- Stream.fold (Fold.partitionByM f Fold.length Fold.length) (Stream.enumerateFromTo (1 :: Int) 100)
--  print r
-- :}
--
-- >>> main
-- (67,33)
--
--
-- This is the consumer side dual of the producer side 'mergeBy' operation.
--
-- When one fold is done, any input meant for it is ignored until the other
-- fold is also done.
--
-- Stops when both the folds stop.
--
-- /See also: 'partitionByFstM' and 'partitionByMinM'./
--
-- /Pre-release/
{-# INLINE partitionByM #-}
partitionByM :: Monad m
    => (a -> m (Either b c)) -> Fold m b x -> Fold m c y -> Fold m a (x, y)
partitionByM = partitionByMUsing teeWith

-- | Similar to 'partitionByM' but terminates when the first fold terminates.
--
--
{-# INLINE partitionByFstM #-}
partitionByFstM :: Monad m
    => (a -> m (Either b c)) -> Fold m b x -> Fold m c y -> Fold m a (x, y)
partitionByFstM = partitionByMUsing teeWithFst

-- | Similar to 'partitionByM' but terminates when any fold terminates.
--
--
{-# INLINE partitionByMinM #-}
partitionByMinM :: Monad m =>
    (a -> m (Either b c)) -> Fold m b x -> Fold m c y -> Fold m a (x, y)
partitionByMinM = partitionByMUsing teeWithMin

-- Note: we could use (a -> Bool) instead of (a -> Either b c), but the latter
-- makes the signature clearer as to which case belongs to which fold.
-- XXX need to check the performance in both cases.
-- | Same as 'partitionByM' but with a pure partition function.
--
-- Count even and odd numbers in a stream:
--
-- >>> :{
--  let f = Fold.partitionBy (\n -> if even n then Left n else Right n)
--                      (fmap (("Even " ++) . show) Fold.length)
--                      (fmap (("Odd "  ++) . show) Fold.length)
--   in Stream.fold f (Stream.enumerateFromTo 1 100)
-- :}
-- ("Even 50","Odd 50")
--
-- /Pre-release/
{-# INLINE partitionBy #-}
partitionBy :: Monad m
    => (a -> Either b c) -> Fold m b x -> Fold m c y -> Fold m a (x, y)
partitionBy f = partitionByM (return . f)

-- | Compose two folds such that the combined fold accepts a stream of 'Either'
-- and routes the 'Left' values to the first fold and 'Right' values to the
-- second fold.
--
-- > partition = partitionBy id
--
-- @since 0.7.0
{-# INLINE partition #-}
partition :: Monad m
    => Fold m b x -> Fold m c y -> Fold m (Either b c) (x, y)
partition = partitionBy id

{-
-- | Send one item to each fold in a round-robin fashion. This is the consumer
-- side dual of producer side 'mergeN' operation.
--
-- partitionN :: Monad m => [Fold m a b] -> Fold m a [b]
-- partitionN fs = Fold step begin done
-}

------------------------------------------------------------------------------
-- demux: in a key value stream fold each key sub-stream with a different fold
------------------------------------------------------------------------------

-- TODO Demultiplex an input element into a number of typed variants. We want
-- to statically restrict the target values within a set of predefined types,
-- an enumeration of a GADT.
--
-- This is the consumer side dual of the producer side 'mux' operation (XXX to
-- be implemented).
--
-- XXX If we use Refold in it, it can perhaps fuse/be more efficient. For
-- example we can store just the result rather than storing the whole fold in
-- the Map.
--
-- Note: There are separate functions to determine Key and Fold from the input
-- because key is to be determined on each input whereas fold is to be
-- determined only once for a key.

-- | In a key value stream, fold values corresponding to each key with a key
-- specific fold. The fold returns the fold result as the second component of
-- the output tuple whenever a fold terminates. The first component of the
-- tuple is a container of in-progress folds. If a fold terminates, another
-- instance of the fold is started upon receiving an input with that key.
--
-- This can be used to scan a stream and collect the results from the scan
-- output.
--
-- /Pre-release/
--
{-# INLINE demuxScanWith #-}
demuxScanWith :: (Monad m, IsMap f, Traversable f) =>
       (a -> Key f)
    -> (a -> m (Fold m a b))
    -> Fold m a (m (f b), Maybe (Key f, b))
demuxScanWith getKey getFold = fmap extract $ foldlM' step initial

    where

    initial = return $ Tuple' IsMap.mapEmpty Nothing

    {-# INLINE runFold #-}
    runFold kv (Fold step1 initial1 extract1) (k, a) = do
         res <- initial1
         case res of
            Partial s -> do
                res1 <- step1 s a
                return
                    $ case res1 of
                        Partial _ ->
                            let fld = Fold step1 (return res1) extract1
                             in Tuple' (IsMap.mapInsert k fld kv) Nothing
                        Done b -> Tuple' (IsMap.mapDelete k kv) (Just (k, b))
            Done b -> return $ Tuple' kv (Just (k, b))

    step (Tuple' kv _) a = do
        let k = getKey a
        case IsMap.mapLookup k kv of
            Nothing -> do
                fld <- getFold a
                runFold kv fld (k, a)
            Just f -> runFold kv f (k, a)

    extract (Tuple' kv x) = (Prelude.mapM f kv, x)

        where

        f (Fold _ i e) = do
            r <- i
            case r of
                Partial s -> e s
                Done b -> return b

-- | This is specialized version of 'demuxScanWith' that uses mutable cells as
-- fold accumulators for better performance.
--
-- >>> demuxScanMutWith = Fold.demuxScanWith
--
{-# INLINE demuxScanMutWith #-}
demuxScanMutWith :: (MonadIO m, IsMap f, Traversable f) =>
       (a -> Key f)
    -> (a -> m (Fold m a b))
    -> Fold m a (m (f b), Maybe (Key f, b))
demuxScanMutWith getKey getFold = fmap extract $ foldlM' step initial

    where

    initial = return $ Tuple' IsMap.mapEmpty Nothing

    {-# INLINE initFold #-}
    initFold kv (Fold step1 initial1 extract1) (k, a) = do
         res <- initial1
         case res of
            Partial s -> do
                res1 <- step1 s a
                case res1 of
                    Partial _ -> do
                        let fld = Fold step1 (return res1) extract1
                        ref <- liftIO $ newIORef fld
                        return $ Tuple' (IsMap.mapInsert k ref kv) Nothing
                    Done b -> return $ Tuple' kv (Just (k, b))
            Done b -> return $ Tuple' kv (Just (k, b))

    {-# INLINE runFold #-}
    runFold kv ref (Fold step1 initial1 extract1) (k, a) = do
         res <- initial1
         case res of
            Partial s -> do
                res1 <- step1 s a
                case res1 of
                        Partial _ -> do
                            let fld = Fold step1 (return res1) extract1
                            liftIO $ writeIORef ref fld
                            return $ Tuple' kv Nothing
                        Done b ->
                            let kv1 = IsMap.mapDelete k kv
                             in return $ Tuple' kv1 (Just (k, b))
            Done _ -> error "demuxScanMutWith: unreachable"

    step (Tuple' kv _) a = do
        let k = getKey a
        case IsMap.mapLookup k kv of
            Nothing -> do
                f <- getFold a
                initFold kv f (k, a)
            Just ref -> do
                f <- liftIO $ readIORef ref
                runFold kv ref f (k, a)

    extract (Tuple' kv x) = (Prelude.mapM f kv, x)

        where

        f ref = do
            (Fold _ i e) <- liftIO $ readIORef ref
            r <- i
            case r of
                Partial s -> e s
                Done b -> return b

-- | This collects all the results of 'demuxScanWith' in a container.
--
{-# INLINE demuxWith #-}
demuxWith :: (Monad m, IsMap f, Traversable f) =>
    (a -> Key f) -> (a -> m (Fold m a b)) -> Fold m a (f b)
demuxWith getKey getFold =
    let
        classifier = demuxScanWith getKey getFold
        getMap Nothing = pure IsMap.mapEmpty
        getMap (Just action) = action
        aggregator =
            teeWith IsMap.mapUnion
                (rmapM getMap $ lmap fst last)
                (lmap snd $ catMaybes toMap)
    in postscan classifier aggregator

-- | Same as 'demuxWith' but uses 'demuxScanMutWith' for better performance.
--
-- >>> demuxMutWith = Fold.demuxWith
--
{-# INLINE demuxMutWith #-}
demuxMutWith :: (MonadIO m, IsMap f, Traversable f) =>
    (a -> Key f) -> (a -> m (Fold m a b)) -> Fold m a (f b)
demuxMutWith getKey getFold =
    let
        classifier = demuxScanMutWith getKey getFold
        getMap Nothing = pure IsMap.mapEmpty
        getMap (Just action) = action
        aggregator =
            teeWith IsMap.mapUnion
                (rmapM getMap $ lmap fst last)
                (lmap snd $ catMaybes toMap)
    in postscan classifier aggregator

-- | Fold a stream of key value pairs using a function that maps keys to folds.
--
-- >>> import Data.Map (Map)
-- >>> :{
--  let f "SUM" = return Fold.sum
--      f _ = return Fold.product
--      input = Stream.fromList [("SUM",1),("PRODUCT",2),("SUM",3),("PRODUCT",4)]
--   in Stream.fold (Fold.demux f) input :: IO (Map String Int)
-- :}
-- fromList [("PRODUCT",8),("SUM",4)]
--
-- >>> demux f = Fold.demuxWith fst (Fold.lmap snd . f)
--
-- /Pre-release/
{-# INLINE demux #-}
demux :: (Monad m, IsMap f, Traversable f) =>
    (Key f -> m (Fold m a b)) -> Fold m (Key f, a) (f b)
demux f = demuxWith fst (\(k, _) -> fmap (lmap snd) (f k))

------------------------------------------------------------------------------
-- Classify: Like demux but uses the same fold for all keys.
------------------------------------------------------------------------------

-- XXX Change these to make the behavior similar to demux* variants. We can
-- implement this using classifyScanManyWith. Maintain a set of done folds in
-- the underlying monad, and when initial is called look it up, if the fold is
-- done then initial would set a flag in the state to ignore the input or
-- return an error.

-- | Folds the values for each key using the supplied fold. When scanning, as
-- soon as the fold is complete, its result is available in the second
-- component of the tuple.  The first component of the tuple is a snapshot of
-- the in-progress folds.
--
-- Once the fold for a key is done, any future values of the key are ignored.
--
-- >>> classifyScanWith f fld = Fold.demuxScanWith f (const fld)
--
{-# INLINE classifyScanWith #-}
classifyScanWith :: (Monad m, IsMap f, Traversable f, Ord (Key f)) =>
    -- Note: we need to return the Map itself to display the in-progress values
    -- e.g. to implement top. We could possibly create a separate abstraction
    -- for that use case. We return an action because we want it to be lazy so
    -- that the downstream consumers can choose to process or discard it.
    (a -> Key f) -> Fold m a b -> Fold m a (m (f b), Maybe (Key f, b))
classifyScanWith f (Fold step1 initial1 extract1) =
    fmap extract $ foldlM' step initial

    where

    initial = return $ Tuple3' IsMap.mapEmpty Set.empty Nothing

    {-# INLINE initFold #-}
    initFold kv set k a = do
        x <- initial1
        case x of
              Partial s -> do
                r <- step1 s a
                return
                    $ case r of
                          Partial s1 ->
                            Tuple3' (IsMap.mapInsert k s1 kv) set Nothing
                          Done b ->
                            Tuple3' kv set (Just (k, b))
              Done b -> return (Tuple3' kv (Set.insert k set) (Just (k, b)))

    step (Tuple3' kv set _) a = do
        let k = f a
        case IsMap.mapLookup k kv of
            Nothing -> do
                if Set.member k set
                then return (Tuple3' kv set Nothing)
                else initFold kv set k a
            Just s -> do
                r <- step1 s a
                return
                    $ case r of
                          Partial s1 ->
                            Tuple3' (IsMap.mapInsert k s1 kv) set Nothing
                          Done b ->
                            let kv1 = IsMap.mapDelete k kv
                             in Tuple3' kv1 (Set.insert k set) (Just (k, b))

    extract (Tuple3' kv _ x) = (Prelude.mapM extract1 kv, x)

-- XXX we can use a Prim IORef if we can constrain the state "s" to be Prim
--
-- The code is almost the same as classifyScanWith except the IORef operations.
--
-- | Same as classifyScanWith except that it uses mutable IORef cells in the
-- Map providing better performance. Be aware that if this is used as a scan,
-- the values in the intermediate Maps would be mutable.
--
-- >>> classifyScanMutWith = classifyScanWith
-- >>> classifyScanMutWith f fld = demuxScanMutWith f (const fld)
--
{-# INLINE classifyScanMutWith #-}
classifyScanMutWith :: (MonadIO m, IsMap f, Traversable f, Ord (Key f)) =>
    (a -> Key f) -> Fold m a b -> Fold m a (m (f b), Maybe (Key f, b))
classifyScanMutWith f (Fold step1 initial1 extract1) =
    fmap extract $ foldlM' step initial

    where

    initial = return $ Tuple3' IsMap.mapEmpty Set.empty Nothing

    {-# INLINE initFold #-}
    initFold kv set k a = do
        x <- initial1
        case x of
              Partial s -> do
                r <- step1 s a
                case r of
                      Partial s1 -> do
                        ref <- liftIO $ newIORef s1
                        return $ Tuple3' (IsMap.mapInsert k ref kv) set Nothing
                      Done b ->
                        return $ Tuple3' kv set (Just (k, b))
              Done b -> return (Tuple3' kv (Set.insert k set) (Just (k, b)))

    step (Tuple3' kv set _) a = do
        let k = f a
        case IsMap.mapLookup k kv of
            Nothing -> do
                if Set.member k set
                then return (Tuple3' kv set Nothing)
                else initFold kv set k a
            Just ref -> do
                s <- liftIO $ readIORef ref
                r <- step1 s a
                case r of
                      Partial s1 -> do
                        liftIO $ writeIORef ref s1
                        return $ Tuple3' kv set Nothing
                      Done b ->
                        let kv1 = IsMap.mapDelete k kv
                         in return
                                $ Tuple3' kv1 (Set.insert k set) (Just (k, b))

    extract (Tuple3' kv _ x) =
        (Prelude.mapM (\ref -> liftIO (readIORef ref) >>= extract1) kv, x)

-- | Fold a key value stream to a key-value container. If the same key appears
-- multiple times, only the last value is retained.
{-# INLINE toMap #-}
toMap :: (Monad m, IsMap f) => Fold m (Key f, a) (f a)
toMap = foldl' (\kv (k, v) -> IsMap.mapInsert k v kv) IsMap.mapEmpty

-- | Split the input stream based on a key field and fold each split using the
-- given fold. Useful for map/reduce, bucketizing the input in different bins
-- or for generating histograms.
--
-- >>> import Data.Map.Strict (Map)
-- >>> :{
--  let input = Stream.fromList [("ONE",1),("ONE",1.1),("TWO",2), ("TWO",2.2)]
--      classify = Fold.classifyWith fst (Fold.lmap snd Fold.toList)
--   in Stream.fold classify input :: IO (Map String [Double])
-- :}
-- fromList [("ONE",[1.0,1.1]),("TWO",[2.0,2.2])]
--
-- Once the classifier fold terminates for a particular key any further inputs
-- in that bucket are ignored.
--
-- Space used is proportional to the number of keys seen till now and
-- monotonically increases because it stores whether a key has been seen or
-- not.
--
-- /Stops: never/
--
-- /Pre-release/
--
{-# INLINE classifyWith #-}
classifyWith :: (Monad m, IsMap f, Traversable f, Ord (Key f)) =>
    (a -> Key f) -> Fold m a b -> Fold m a (f b)
classifyWith f fld =
    let
        classifier = classifyScanWith f fld
        getMap Nothing = pure IsMap.mapEmpty
        getMap (Just action) = action
        aggregator =
            teeWith IsMap.mapUnion
                (rmapM getMap $ lmap fst last)
                (lmap snd $ catMaybes toMap)
    in postscan classifier aggregator

-- | Same as 'classifyWith' but maybe faster because it uses mutable cells as
-- fold accumulators in the Map.
--
-- >>> classifyMutWith = Fold.classifyWith
--
{-# INLINE classifyMutWith #-}
classifyMutWith :: (MonadIO m, IsMap f, Traversable f, Ord (Key f)) =>
    (a -> Key f) -> Fold m a b -> Fold m a (f b)
classifyMutWith f fld =
    let
        classifier = classifyScanMutWith f fld
        getMap Nothing = pure IsMap.mapEmpty
        getMap (Just action) = action
        aggregator =
            teeWith IsMap.mapUnion
                (rmapM getMap $ lmap fst last)
                (lmap snd $ catMaybes toMap)
    in postscan classifier aggregator

-- | Given an input stream of key value pairs and a fold for values, fold all
-- the values belonging to each key.  Useful for map/reduce, bucketizing the
-- input in different bins or for generating histograms.
--
-- >>> :{
--  let input = Stream.fromList [("ONE",1),("ONE",1.1),("TWO",2), ("TWO",2.2)]
--   in Stream.fold (Fold.classify Fold.toList) input
-- :}
-- fromList [("ONE",[1.0,1.1]),("TWO",[2.0,2.2])]
--
-- Same as:
--
-- >>> classify = Fold.classifyWith fst . Fold.lmap snd
--
-- /Pre-release/
{-# INLINE classify #-}
classify :: (Monad m, Ord k) => Fold m a b -> Fold m (k, a) (Map k b)
classify = classifyWith fst . lmap snd

-- | Determine the frequency of each element in the stream.
--
-- You can just collect the keys of the resulting map to get the unique
-- elements in the stream.
--
{-# INLINE frequency #-}
frequency :: (Monad m, Ord a) => Fold m a (Map a Int)
frequency = classifyWith id length

------------------------------------------------------------------------------
-- Unzipping
------------------------------------------------------------------------------

{-# INLINE unzipWithMUsing #-}
unzipWithMUsing :: Monad m =>
       (  (x -> y -> (x, y))
       -> Fold m (b, c) x
       -> Fold m (b, c) y
       -> Fold m (b, c) (x, y)
       )
    -> (a -> m (b, c))
    -> Fold m b x
    -> Fold m c y
    -> Fold m a (x, y)
unzipWithMUsing t f fld1 fld2 =
    let f1 = lmap fst fld1  -- :: Fold m (b, c) b
        f2 = lmap snd fld2  -- :: Fold m (b, c) c
     in lmapM f (t (,) f1 f2)

-- | Like 'unzipWith' but with a monadic splitter function.
--
-- @unzipWithM k f1 f2 = lmapM k (unzip f1 f2)@
--
-- /Pre-release/
{-# INLINE unzipWithM #-}
unzipWithM :: Monad m
    => (a -> m (b,c)) -> Fold m b x -> Fold m c y -> Fold m a (x,y)
unzipWithM = unzipWithMUsing teeWith

-- | Similar to 'unzipWithM' but terminates when the first fold terminates.
--
--
{-# INLINE unzipWithFstM #-}
unzipWithFstM :: Monad m =>
    (a -> m (b, c)) -> Fold m b x -> Fold m c y -> Fold m a (x, y)
unzipWithFstM = unzipWithMUsing teeWithFst

-- | Similar to 'unzipWithM' but terminates when any fold terminates.
--
--
{-# INLINE unzipWithMinM #-}
unzipWithMinM :: Monad m =>
    (a -> m (b,c)) -> Fold m b x -> Fold m c y -> Fold m a (x,y)
unzipWithMinM = unzipWithMUsing teeWithMin

-- | Split elements in the input stream into two parts using a pure splitter
-- function, direct each part to a different fold and zip the results.
--
-- @unzipWith f fld1 fld2 = Fold.lmap f (Fold.unzip fld1 fld2)@
--
-- This fold terminates when both the input folds terminate.
--
-- /Pre-release/
{-# INLINE unzipWith #-}
unzipWith :: Monad m
    => (a -> (b,c)) -> Fold m b x -> Fold m c y -> Fold m a (x,y)
unzipWith f = unzipWithM (return . f)

-- | Send the elements of tuples in a stream of tuples through two different
-- folds.
--
-- @
--
--                           |-------Fold m a x--------|
-- ---------stream of (a,b)--|                         |----m (x,y)
--                           |-------Fold m b y--------|
--
-- @
--
-- > unzip = Fold.unzipWith id
--
-- This is the consumer side dual of the producer side 'zip' operation.
--
-- @since 0.7.0
{-# INLINE unzip #-}
unzip :: Monad m => Fold m a x -> Fold m b y -> Fold m (a,b) (x,y)
unzip = unzipWith id

------------------------------------------------------------------------------
-- Combining streams and folds - Zipping
------------------------------------------------------------------------------

-- XXX These can be implemented using the fold scan, using the stream as a
-- state.
-- XXX Stream Skip state cannot be efficiently handled in folds but can be
-- handled in parsers using the Continue facility. See zipWithM in the Parser
-- module.
--
-- | Zip a stream with the input of a fold using the supplied function.
--
-- /Unimplemented/
--
{-# INLINE zipWithM #-}
zipWithM :: -- Monad m =>
    (a -> b -> m c) -> t m a -> Fold m c x -> Fold m b x
zipWithM = undefined

-- | Zip a stream with the input of a fold.
--
-- /Unimplemented/
--
{-# INLINE zip #-}
zip :: Monad m => t m a -> Fold m (a, b) x -> Fold m b x
zip = zipWithM (curry return)

-- | Pair each element of a fold input with its index, starting from index 0.
--
{-# INLINE indexingWith #-}
indexingWith :: Monad m => Int -> (Int -> Int) -> Fold m a (Maybe (Int, a))
indexingWith i f = fmap toMaybe $ foldl' step initial

    where

    initial = Nothing'

    step Nothing' a = Just' (i, a)
    step (Just' (n, _)) a = Just' (f n, a)

{-# INLINE indexing #-}
indexing :: Monad m => Fold m a (Maybe (Int, a))
indexing = indexingWith 0 (+ 1)

{-# INLINE indexingRev #-}
indexingRev :: Monad m => Int -> Fold m a (Maybe (Int, a))
indexingRev n = indexingWith n (subtract 1)

-- | Pair each element of a fold input with its index, starting from index 0.
--
{-# INLINE indexed #-}
indexed :: Monad m => Fold m (Int, a) b -> Fold m a b
indexed = scanMaybe indexing

-- | Change the predicate function of a Fold from @a -> b@ to accept an
-- additional state input @(s, a) -> b@. Convenient to filter with an
-- addiitonal index or time input.
--
-- @
-- filterWithIndex = with indexed filter
-- filterWithAbsTime = with timestamped filter
-- filterWithRelTime = with timeIndexed filter
-- @
--
-- /Pre-release/
{-# INLINE with #-}
with ::
       (Fold m (s, a) b -> Fold m a b)
    -> (((s, a) -> c) -> Fold m (s, a) b -> Fold m (s, a) b)
    -> (((s, a) -> c) -> Fold m a b -> Fold m a b)
with f comb g = f . comb g . lmap snd

-- XXX Implement as a filter
-- sampleFromthen :: Monad m => Int -> Int -> Fold m a (Maybe a)

-- | @sampleFromthen offset stride@ samples the element at @offset@ index and
-- then every element at strides of @stride@.
--
{-# INLINE sampleFromthen #-}
sampleFromthen :: Monad m => Int -> Int -> Fold m a b -> Fold m a b
sampleFromthen offset size =
    with indexed filter (\(i, _) -> (i + offset) `mod` size == 0)

------------------------------------------------------------------------------
-- Nesting
------------------------------------------------------------------------------

-- | @concatSequence f t@ applies folds from stream @t@ sequentially and
-- collects the results using the fold @f@.
--
-- /Unimplemented/
--
{-# INLINE concatSequence #-}
concatSequence ::
    -- IsStream t =>
    Fold m b c -> t (Fold m a b) -> Fold m a c
concatSequence _f _p = undefined

-- | Group the input stream into groups of elements between @low@ and @high@.
-- Collection starts in chunks of @low@ and then keeps doubling until we reach
-- @high@. Each chunk is folded using the provided fold function.
--
-- This could be useful, for example, when we are folding a stream of unknown
-- size to a stream of arrays and we want to minimize the number of
-- allocations.
--
-- NOTE: this would be an application of "many" using a terminating fold.
--
-- /Unimplemented/
--
{-# INLINE chunksBetween #-}
chunksBetween :: -- Monad m =>
       Int -> Int -> Fold m a b -> Fold m b c -> Fold m a c
chunksBetween _low _high _f1 _f2 = undefined

-- | A fold that buffers its input to a pure stream.
--
-- /Warning!/ working on large streams accumulated as buffers in memory could
-- be very inefficient, consider using "Streamly.Data.Array" instead.
--
-- >>> toStream = fmap Stream.fromStreamK Fold.toStreamK
--
-- /Pre-release/
{-# INLINE toStream #-}
toStream :: Monad m => Fold m a (Stream n a)
toStream = fmap Stream.fromStreamK toStreamK

-- This is more efficient than 'toStream'. toStream is exactly the same as
-- reversing the stream after toStreamRev.
--
-- | Buffers the input stream to a pure stream in the reverse order of the
-- input.
--
-- >>> toStreamRev = fmap Stream.fromStreamK Fold.toStreamKRev
--
-- /Warning!/ working on large streams accumulated as buffers in memory could
-- be very inefficient, consider using "Streamly.Data.Array" instead.
--
-- /Pre-release/

--  xn : ... : x2 : x1 : []
{-# INLINE toStreamRev #-}
toStreamRev :: Monad m => Fold m a (Stream n a)
toStreamRev = fmap Stream.fromStreamK toStreamKRev

-- XXX This does not fuse. It contains a recursive step function. We will need
-- a Skip input constructor in the fold type to make it fuse.
--
-- | Unfold and flatten the input stream of a fold.
--
-- @
-- Stream.fold (unfoldMany u f) = Stream.fold f . Stream.unfoldMany u
-- @
--
-- /Pre-release/
{-# INLINE unfoldMany #-}
unfoldMany :: Monad m => Unfold m a b -> Fold m b c -> Fold m a c
unfoldMany (Unfold ustep inject) (Fold fstep initial extract) =
    Fold consume initial extract

    where

    {-# INLINE produce #-}
    produce fs us = do
        ures <- ustep us
        case ures of
            StreamD.Yield b us1 -> do
                fres <- fstep fs b
                case fres of
                    Partial fs1 -> produce fs1 us1
                    -- XXX What to do with the remaining stream?
                    Done c -> return $ Done c
            StreamD.Skip us1 -> produce fs us1
            StreamD.Stop -> return $ Partial fs

    {-# INLINE_LATE consume #-}
    consume s a = inject a >>= produce s

{-# INLINE topBy #-}
topBy :: (MonadIO m, Unboxed a) =>
       (a -> a -> Ordering)
    -> Int
    -> Fold m a (MA.Array a)
topBy cmp n = Fold step initial extract

    where

    initial = do
        arr <- MA.newPinned n
        if n <= 0
        then return $ Done arr
        else return $ Partial (arr, 0)

    step (arr, i) x =
        if i < n
        then do
            arr' <- MA.snoc arr x
            MA.bubble cmp arr'
            return $ Partial (arr', i + 1)
        else do
            x1 <- MA.getIndexUnsafe (i - 1) arr
            case x `cmp` x1 of
                LT -> do
                    MA.putIndexUnsafe (i - 1) x arr
                    MA.bubble cmp arr
                    return $ Partial (arr, i)
                _ -> return $ Partial (arr, i)

    extract = return . fst

-- | Fold the input stream to top n elements.
--
-- >>> stream = Stream.fromList [2::Int,7,9,3,1,5,6,11,17]
-- >>> Stream.fold (Fold.top 3) stream >>= MA.toList
-- [17,11,9]
--
-- /Pre-release/
{-# INLINE top #-}
top :: (MonadIO m, Unboxed a, Ord a) => Int -> Fold m a (MA.Array a)
top = topBy $ flip compare

-- | Fold the input stream to bottom n elements.
--
-- >>> stream = Stream.fromList [2::Int,7,9,3,1,5,6,11,17]
-- >>> Stream.fold (Fold.bottom 3) stream >>= MA.toList
-- [1,2,3]
--
-- /Pre-release/
{-# INLINE bottom #-}
bottom :: (MonadIO m, Unboxed a, Ord a) => Int -> Fold m a (MA.Array a)
bottom = topBy compare