-- SPDX-License-Identifier: GPL-3.0-or-later
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DeriveTraversable #-}
module Zerem
  (
  -- * Purpose
  -- $purpose
  -- * How to write a stream.
  -- $how-to-write-a-stream
    Zerem(..)
  , Step(..)
  -- , mapState
  -- , mapYield
  -- , mapYieldM
  -- , mapYieldF
  -- , mapYieldFM
  , next
  , next_
  , InitZerem(..)
  , initZerem
  , ZipZerem(..)
  , ListZerem(..)
  , Pipe(..)
  -- * Mapping
  , map
  , maps
  , mapM
  , mapsM
  , flatLift
  , flatMapLift
  , (>>^)
  , hoistAny
  -- * Folding
  , runZerem
  , runZerem_
  , map_
  , maps_
  , mapM_
  , mapsM_
  , Zerem.foldr
  , foldr_
  , Zerem.foldrM
  , foldrM_
  , Zerem.foldl
  , foldl_
  , Zerem.foldlM
  , foldlM_
  -- * Scanning
  , Zerem.scanl
  , scanl_
  , scansl
  , scansl_
  -- * Filtering
  , scanMaybel
  , scanMaybel_
  , catMaybes
  , mapMaybe
  , Zerem.concat
  , Zerem.concatMap
  , Zerem.filter
  , Zerem.take
  , Zerem.drop
  -- * Generating
  , yield
  , yieldMany
  , yieldM
  , yieldManyM
  , concatStream
  , mtimes
  ) where

import Prelude hiding (id, (.), map, mapM, mapM_)
import Control.Category
-- import Data.Profunctor
import Control.Applicative
import Data.Functor
import Data.Functor.Identity
import Data.Foldable hiding (mapM_)
import Data.Void
import Data.Bifunctor
import Data.Bitraversable
import Control.Monad hiding (mapM, mapM_)
import Data.Semigroup ( Semigroup(..) )
import Data.List.NonEmpty (NonEmpty(..))
import Control.Exception ( assert )
import Control.Monad.Morph ( MonadTrans(..), MFunctor(..) )
import GHC.Exts ( coerce, Coercible )
import Numeric.Natural ( Natural )
import Control.Monad.Logic.Class ( MonadLogic(..) )
import Control.Comonad ( Comonad(..), ComonadApply(..) )
import Data.Profunctor
import Data.Profunctor.Monad

{-$how-to-write-a-stream
.
-}

( #. ) :: Coercible c b => (b -> c) -> (a -> b) -> (a -> c)
( #. ) _ = coerce (\x -> x :: b) :: forall a b. Coercible b a => a -> b
{-# INLINE ( #. ) #-}

(.#) :: Coercible b a => (b -> c) -> (a -> b) -> (a -> c)
(.#) pbc _ = coerce pbc
{-# INLINE (.#) #-}

infixr 9 #.
infixl 8 .#

data Zerem f m r = forall s. Zerem (s -> m (Step s f r)) (m s)
{-^
The stream datatype. It consists of a state machine, and an action
that yields the initial state.
-}

data Step s f r
  = Done r -- ^ Stream is over and @r@ is the result.
  | Skip s -- ^ Stream is transitioning state without yielding.
  | Yield !(f s) -- ^ Stream yields a functor. Usually @(a, s)@.
  deriving (Functor, Foldable, Traversable)

type ZeremOf a = Zerem ((,) a)
type Stepper s f m u = s -> m (Step s f u)

mapState :: Functor f => (s -> s') -> Step s f u -> Step s' f u
mapState f = \case
  Done u -> Done u
  Skip s -> Skip (f s)
  Yield fs -> Yield (f <$> fs)
{-# INLINE mapState #-}

mapYield :: (f s -> g s) -> Step s f u -> Step s g u
mapYield phi = \case
  Done u -> Done u
  Skip s -> Skip s
  Yield fs -> Yield $ phi fs
{-# INLINE mapYield #-}

mapYieldM :: Monad m => (f s -> m (g s)) -> Step s f u -> m (Step s g u)
mapYieldM phi = \case
  Done u -> pure $ Done u
  Skip s -> pure $ Skip s
  Yield fs -> Yield <$> phi fs
{-# INLINE mapYieldM #-}

mapYieldF :: (forall x. f x -> g x) -> Step s f u -> Step s g u
mapYieldF phi = \case
  Done u -> Done u
  Skip s -> Skip s
  Yield fs -> Yield $ phi fs
{-# INLINE mapYieldF #-}

mapYieldFM :: Monad m => (forall x. f x -> m (g x)) -> Step s f u -> m (Step s g u)
mapYieldFM phi = \case
  Done u -> pure $ Done u
  Skip s -> pure $ Skip s
  Yield fs -> Yield <$> phi fs
{-# INLINE mapYieldFM #-}

next :: (Functor f, Monad m) => Zerem f m r -> m (Either r (f (Zerem f m r)))
next (Zerem next s) = loop =<< s where
  loop s =
    next s >>= \case
      Done r -> pure $ Left r
      Skip s' -> loop s'
      Yield fs' -> pure . Right $ fmap (Zerem next . pure) fs'

next_ :: (Functor f, Monad m) => Zerem f m r -> m (Maybe (f (Zerem f m r)))
next_ (Zerem next s) = loop =<< s where
  loop s =
    next s >>= \case
      Done _ -> pure Nothing
      Skip s' -> loop s'
      Yield fs' -> pure . Just $ fmap (Zerem next . pure) fs'

deriving instance Functor m => Functor (Zerem f m)

{-|
Not lazy on its arguments. If you use something like:

@
as = yield 'a' *> as
@

the program will go into an infinite loop.

For that kind of construction, consider using functions
from the `Generating` section, or simply writing a stream
by hand:

@
as = Zerem next (pure a) where
  next a = pure $ Z.Yield (a, a)
@
-}
instance (Functor f, Monad m) => Applicative (Zerem f m) where
  pure r = Zerem (pure . Done) (pure r)

  (Zerem nextf sf) <*> (Zerem nexta sa) = Zerem next' (Left <$> sf) where
    next' (Left s) =
      nextf s >>= \case
        Done f -> sa <&> \sa -> Skip (Right (f, sa))
        Skip s' -> pure $ Skip (Left s')
        Yield fs' -> pure $ Yield (fmap Left fs')
    next' (Right (f, s)) =
      nexta s <&> \case
        Done a -> Done $ f a
        Skip s' -> Skip (Right (f, s'))
        Yield fs' -> Yield (fmap (Right . (f,)) fs')

  liftA2 f (Zerem nexta sa) (Zerem nextb sb) = Zerem next' (Left <$> sa) where
    next' (Left s) =
      nexta s >>= \case
        Done a -> sb <&> \sb -> Skip (Right (a, sb))
        Skip s' -> pure $ Skip (Left s')
        Yield fs' -> pure $ Yield (fmap Left fs')
    next' (Right (a, s)) =
      nextb s <&> \case
        Done b -> Done (f a b)
        Skip s' -> Skip (Right (a, s'))
        Yield fs' -> Yield (fmap (Right . (a,)) fs')

  (Zerem nextf sf) *> (Zerem nexta sa) = Zerem next' (Left <$> sf) where
    next' (Left sf) =
      nextf sf >>= \case
        Done _ -> sa <&> \sa -> Skip (Right sa)
        Skip s' -> pure $ Skip (Left s')
        Yield fs' -> pure $ Yield (fmap Left fs')
    next' (Right sa) =
      nexta sa <&> \case
        Done u -> Done u
        Skip s' -> Skip (Right s')
        Yield fs' -> Yield (fmap Right fs')

  {-# INLINE pure #-}
  {-# INLINE (<*>) #-}
  {-# INLINE liftA2 #-}

data InitZerem f m r = forall s. InitZerem (s -> m (Step s f r)) s
-- deriving instance Functor m => Functor (InitZerem f m)

initZerem :: Monad m => Zerem f m r -> m (InitZerem f m r)
initZerem (Zerem next ms) = ms <&> \s -> InitZerem next s

instance (Functor f, Monad m) => Monad (Zerem f m) where
  (Zerem nexta sa) >>= k = Zerem next' (Left <$> sa) where
    next' (Left s) =
      nexta s >>= \case
        -- Bye bye fusion ☹️
        Done a -> Skip . Right <$> initZerem (k a)
        Skip s' -> pure $ Skip (Left s')
        Yield fs' -> pure $ Yield (fmap Left fs')
    next' (Right (InitZerem next s)) =
      next s <&> \case
        Done r -> Done r
        Skip s' -> Skip (Right $ InitZerem next s')
        Yield fs' -> Yield $ fmap (Right . InitZerem next) fs'

  (>>) = (*>)

  {-# INLINE (>>=) #-}

instance MonadTrans (Zerem a) where
  lift = Zerem (pure . Done)

zliftA :: Applicative m => m r -> Zerem a m r
zliftA = Zerem (pure . Done)

instance (Functor f, Monad m) => Semigroup (Zerem f m r) where
  (Zerem nextl sl) <> (Zerem nextr sr) = Zerem next' (liftA2 (curry Left) sl sr) where
    next' (Left (sl, sr)) =
      nextl sl <&> \case
        Done _ -> Skip (Right sr)
        Skip s' -> Skip (Left (s', sr))
        Yield fs' -> Yield $ fmap (\s' -> Left (s', sr)) fs'
    next' (Right sr) =
      nextr sr <&> \case
        Done u -> Done u
        Skip s' -> Skip (Right s')
        Yield fs' -> Yield $ fmap Right fs'

  sconcat nonempty = Zerem next' (pure (Left nonempty)) where
    next' (Left (z :| rest)) =
      (\z -> Skip (Right (z, rest))) <$> initZerem z
    next' (Right (InitZerem next s, rest)) =
      next s >>= \case
        Done u -> case rest of
          [] -> pure $ Done u
          (z : zs) -> (\z -> Skip (Right (z, zs))) <$> initZerem z
        Skip s' -> pure $ Skip (Right (InitZerem next s', rest))
        Yield fs' -> pure . Yield $ fmap (\s' -> Right (InitZerem next s', rest)) fs'

  stimes n (Zerem next s0) = assert (n >= 1) $ Zerem next' ((n,) <$> s0) where
    next' (n, s) =
      next s >>= \case
        -- Using Monad constraint to re-run the initialization action
        -- which might be creating a mutable reference that we don't want
        -- to re-use. Also, it's good to behave like (sconcat . replicate n).
        Done u -> if n == 1 then pure $ Done u else Skip . (pred n,) <$> s0
        Skip s' -> pure $ Skip (n, s')
        Yield fs' -> pure . Yield $ fmap (n,) fs'

instance (Functor f, Monad m) => Monoid (Zerem f m ()) where
  mempty = Zerem (pure . Done) (pure ())

  mconcat list = Zerem next' (pure (Left list)) where
    next' (Left []) = pure (Done ())
    next' (Left (z : rest)) =
      -- Using Monad constraint to run initialization actions.
      Skip . Right . (, rest) <$> initZerem z
    next' (Right (InitZerem next s, rest)) =
      next s <&> \case
        Done u -> Done u
        Skip s' -> Skip (Right (InitZerem next s', rest))
        Yield fs' -> Yield $ fmap (\s' -> Right (InitZerem next s', rest)) fs'

instance MFunctor (Zerem f) where
  hoist f (Zerem next s) = Zerem (f . next) (f s)

newtype ZipZerem r m a = ZipZerem { runZipZerem :: Zerem ((,) a) m r }

instance Functor m => Functor (ZipZerem r m) where
  fmap f (ZipZerem z) = ZipZerem $ map f z

instance Monad m => Applicative (ZipZerem r m) where
  pure a = ZipZerem (Zerem (\a -> pure $ Yield (a, a)) (pure a))

  (ZipZerem (Zerem nextf msf)) <*> (ZipZerem (Zerem nexta msa))
    = ZipZerem (Zerem next' ((\sf -> Left (sf, msa)) <$> msf)) where
    next' (Left (sf, msa)) =
      nextf sf >>= \case
        Done r -> pure $ Done r
        Skip s' -> pure . Skip $ Left (s', msa)
        Yield (f, s') -> msa <&> \sa -> Skip $ Right (f, s', sa)
    next' (Right (f, sf, sa)) =
      nexta sa <&> \case
        Done r -> Done r
        Skip s' -> Skip $ Right (f, sf, s')
        Yield (a, s') -> Yield (f a, Left (sf, pure s'))

  liftA2 f (ZipZerem (Zerem nexta msa)) (ZipZerem (Zerem nextb msb))
    = ZipZerem (Zerem next' ((\sa -> Left (sa, msb)) <$> msa)) where
    next' (Left (sa, msb)) =
      nexta sa >>= \case
        Done r -> pure $ Done r
        Skip s' -> pure . Skip $ Left (s', msb)
        Yield (a, s') -> msb <&> \sb -> Skip $ Right (a, s', sb)
    next' (Right (a, sa, sb)) =
      nextb sb <&> \case
        Done r -> Done r
        Skip s' -> Skip $ Right (a, sa, s')
        Yield (b, s') -> Yield (f a b, Left (sa, pure s'))

  (ZipZerem (Zerem next_ ms_)) *> (ZipZerem (Zerem nexta msa))
    = ZipZerem (Zerem next' ((\s_ -> Left (s_, msa)) <$> ms_)) where
    next' (Left (s_, msa)) =
      next_ s_ >>= \case
        Done r -> pure $ Done r
        Skip s' -> pure . Skip $ Left (s', msa)
        Yield (_, s') -> msa <&> \sa -> Skip $ Right (s', sa)
    next' (Right (s_, sa)) =
      nexta sa <&> \case
        Done r -> Done r
        Skip s' -> Skip $ Right (s_, s')
        Yield (a, s') -> Yield (a, Left (s_, pure s'))

  (ZipZerem (Zerem nexta msa)) <* (ZipZerem (Zerem next_ ms_))
    = ZipZerem (Zerem next' ((\sa -> Left (sa, ms_)) <$> msa)) where
    next' (Left (sa, ms_)) =
      nexta sa >>= \case
        Done r -> pure $ Done r
        Skip s' -> pure . Skip $ Left (s', ms_)
        Yield (a, s') -> ms_ <&> \s_ -> Skip $ Right (a, s', s_)
    next' (Right (a, sa, s_)) =
      next_ s_ <&> \case
        Done r -> Done r
        Skip s' -> Skip $ Right (a, sa, s')
        Yield (_, s') -> Yield (a, Left (sa, pure s'))

instance Monad m => Alternative (ZipZerem () m) where
  empty = ZipZerem $ Zerem (const (pure $ Done ())) (pure ())

  (ZipZerem (Zerem nextl msl)) <|> (ZipZerem (Zerem nextr msr))
    = ZipZerem $ Zerem next (Left . (0 :: Natural,) <$> msl) where
    next (Left (n, sl)) =
      nextl sl >>= \case
        Done () -> msr <&> \sr -> Skip $ Right (n, sr)
        Skip s' -> pure $ Skip $ Left (n, s')
        Yield (a, s') -> pure $ Yield (a, Left (n+1, s'))
    next (Right (0, sr)) =
      nextr sr <&> \case
        Done () -> Done ()
        Skip s' -> Skip $ Right (0, s')
        Yield (a, s') -> Yield (a, Right (0, s'))
    next (Right (n, sr)) =
      nextr sr <&> \case
        Done () -> Done ()
        Skip s' -> Skip $ Right (n, s')
        Yield (_, s') -> Skip $ Right (n-1, s')

instance {-# OVERLAPPING #-} Applicative (ZipZerem Void Identity) where
  pure a = ZipZerem (Zerem (\a -> pure $ Yield (a, a)) (pure a))

  (ZipZerem (Zerem nextf msf)) <*> (ZipZerem (Zerem nexta msa))
    = ZipZerem (Zerem next' ((\sf -> Left (sf, msa)) <$> msf)) where
    next' (Left (sf, msa)) =
      nextf sf >>= \case
        Done r -> pure $ Done r
        Skip s' -> pure . Skip $ Left (s', msa)
        Yield (f, s') -> msa <&> \sa -> Skip $ Right (f, s', sa)
    next' (Right (f, sf, sa)) =
      nexta sa <&> \case
        Done r -> Done r
        Skip s' -> Skip $ Right (f, sf, s')
        Yield (a, s') -> Yield (f a, Left (sf, pure s'))

  liftA2 f (ZipZerem (Zerem nexta msa)) (ZipZerem (Zerem nextb msb))
    = ZipZerem (Zerem next' ((\sa -> Left (sa, msb)) <$> msa)) where
    next' (Left (sa, msb)) =
      nexta sa >>= \case
        Done r -> pure $ Done r
        Skip s' -> pure . Skip $ Left (s', msb)
        Yield (a, s') -> msb <&> \sb -> Skip $ Right (a, s', sb)
    next' (Right (a, sa, sb)) =
      nextb sb <&> \case
        Done r -> Done r
        Skip s' -> Skip $ Right (a, sa, s')
        Yield (b, s') -> Yield (f a b, Left (sa, pure s'))

  _ *> z = z

  z <* _ = z

instance Monad (ZipZerem Void Identity) where
  (ZipZerem (Zerem nexta msa)) >>= k = ZipZerem (Zerem next (Left . (0 :: Natural,) <$> msa)) where
    next (Left (n, sa)) =
      nexta sa >>= \case
        Done v -> absurd v
        Skip s' -> pure . Skip $ Left (n, s')
        Yield (a, s') -> case k a of -- RIP Fusion
          ZipZerem zb -> (\zb -> Skip $ Right (n+1, s', n, zb)) <$> initZerem zb
    next (Right (n, sa, c, InitZerem nextb sb)) =
      nextb sb <&> \case
        Done v -> absurd v
        Skip s' -> Skip $ Right (n, sa, c, InitZerem nextb s')
        Yield (b, s') -> case c of
          0 -> Yield (b, Left (n, sa))
          _ -> Skip $ Right (n, sa, c-1, InitZerem nextb s')

instance MonadTrans (ZipZerem Void) where
  lift ma = ZipZerem (Zerem (\a -> pure $ Yield (a, a)) ma)

instance MFunctor (ZipZerem r) where
  hoist f (ZipZerem (Zerem next s)) = ZipZerem $ Zerem (f . next) (f s)

instance Comonad (ZipZerem Void Identity) where
  extract (ZipZerem (Zerem next ms)) = extract $ loop =<< ms where
    loop s =
      next s >>= \case
        Done v -> absurd v
        Skip s' -> loop s'
        Yield (a, _) -> pure a

  duplicate (ZipZerem z@(Zerem next ms)) = ZipZerem $ Zerem next' ((0 :: Natural,) <$> ms) where
    next' (n, s) =
      next s <&> \case
        Done v -> absurd v
        Skip s' -> Skip (n, s')
        Yield (_, s') -> Yield (ZipZerem $ Zerem.drop (fromIntegral n) z, (n+1, s'))

instance ComonadApply (ZipZerem Void Identity) where
  (<@>) = (<*>)

instance Foldable (ZipZerem r Identity) where
  foldr f z (ZipZerem stream) = runIdentity $ foldr_ f z stream
  {-# INLINE foldr #-}

  foldl' f z (ZipZerem stream) = runIdentity $ foldl_ f z stream
  {-# INLINE foldl' #-}

instance Traversable (ZipZerem () Identity) where
  traverse f = fmap (ZipZerem . yieldMany) . traverse f . toList
  {-# INLINE traverse #-}

newtype ListZerem r m a = ListZerem { runListZerem :: Zerem ((,) a) m r }

instance Functor m => Functor (ListZerem r m) where
  fmap f (ListZerem z) = ListZerem $ map f z

instance Monad m => Applicative (ListZerem () m) where
  pure a = ListZerem (Zerem (pure . next) (pure (Just a))) where
    next Nothing = Done ()
    next (Just a) = Yield (a, Nothing)

  (ListZerem (Zerem nextf msf)) <*> (ListZerem (Zerem nexta msa))
    = ListZerem (Zerem next' (Left <$> msf)) where
    next' (Left sf) =
      nextf sf >>= \case
        Done () -> pure $ Done ()
        Skip s' -> pure . Skip $ Left s'
        Yield (f, s') -> msa <&> \sa -> Skip $ Right (f, s', sa)
    next' (Right (f, sf, sa)) =
      nexta sa <&> \case
        Done _ -> Skip $ Left sf
        Skip s' -> Skip $ Right (f, sf, s')
        Yield (a, s') -> Yield (f a, Right (f, sf, s'))

  liftA2 f (ListZerem (Zerem nexta msa)) (ListZerem (Zerem nextb msb))
    = ListZerem (Zerem next' (Left <$> msa)) where
    next' (Left sa) =
      nexta sa >>= \case
        Done () -> pure $ Done ()
        Skip s' -> pure . Skip $ Left s'
        Yield (a, s') -> msb <&> \sb -> Skip $ Right (a, s', sb)
    next' (Right (a, sa, sb)) =
      nextb sb <&> \case
        Done () -> Skip $ Left sa
        Skip s' -> Skip $ Right (a, sa, s')
        Yield (b, s') -> Yield (f a b, Right (a, sa, s'))

  (ListZerem (Zerem next_ ms_)) *> (ListZerem (Zerem nexta msa))
    = ListZerem (Zerem next' (Left <$> ms_)) where
    next' (Left s_) =
      next_ s_ >>= \case
        Done () -> pure $ Done ()
        Skip s' -> pure . Skip $ Left s'
        Yield (_, s') -> msa <&> \sa -> Skip $ Right (s', sa)
    next' (Right (s_, sa)) =
      nexta sa <&> \case
        Done () -> Skip $ Left s_
        Skip s' -> Skip $ Right (s_, s')
        Yield (a, s') -> Yield (a, Right (s_, s'))

  (ListZerem (Zerem nexta msa)) <* (ListZerem (Zerem next_ ms_))
    = ListZerem (Zerem next' (Left <$> msa)) where
    next' (Left sa) =
      nexta sa >>= \case
        Done r -> pure $ Done r
        Skip s' -> pure . Skip $ Left s'
        Yield (a, s') -> ms_ <&> \s_ -> Skip $ Right (a, s', s_)
    next' (Right (a, sa, s_)) =
      next_ s_ <&> \case
        Done r -> Done r
        Skip s' -> Skip $ Right (a, sa, s')
        Yield (_, s') -> Yield (a, Right (a, sa, s'))

instance Monad m => Alternative (ListZerem () m) where
  empty = ListZerem $ Zerem (const (pure $ Done ())) (pure ())
  (ListZerem zl) <|> (ListZerem zr) = ListZerem $ zl *> zr

instance Monad m => Monad (ListZerem () m) where
  (ListZerem (Zerem nexta msa)) >>= k = ListZerem (Zerem next (Left <$> msa)) where
    next (Left sa) =
      nexta sa >>= \case
        Done () -> pure $ Done ()
        Skip s' -> pure . Skip $ Left s'
        Yield (a, s') -> case k a of -- RIP Fusion
          ListZerem zb -> (\zb -> Skip $ Right (s', zb)) <$> initZerem zb
    next (Right (sa, InitZerem nextb sb)) =
      nextb sb <&> \case
        Done () -> Skip $ Left sa
        Skip s' -> Skip $ Right (sa, InitZerem nextb s')
        Yield (b, s') -> Yield (b, Right (sa, InitZerem nextb s'))

  (>>) = (*>)

instance Monad m => MonadPlus (ListZerem () m)

instance Monad m => MonadLogic (ListZerem () m) where
  msplit (ListZerem (Zerem next s)) = lift $ loop =<< s where
    loop s =
      next s >>= \case
        Done () -> pure Nothing
        Skip s' -> loop s'
        Yield fs' -> pure . Just $
          fmap (ListZerem . Zerem next . pure) fs'
  
  interleave (ListZerem (Zerem nextl msl)) (ListZerem (Zerem nextr msr))
    = ListZerem $ Zerem next ((\sl -> Left (sl, Just msr)) <$> msl) where
    next (Left (sl, Just msr)) =
      nextl sl >>= \case
       Done () -> msr <&> \sr -> Skip $ Right (Nothing, sr)
       Skip s' -> pure . Skip $ Left (s', Just msr)
       Yield (a, s') -> msr <&> \sr -> Yield (a, Right (Just s', sr))
    next (Left (sl, Nothing)) =
      nextl sl <&> \case
       Done () -> Done ()
       Skip s' -> Skip $ Left (s', Nothing)
       Yield (a, s') -> Yield (a, Left (s', Nothing))
    next (Right (Just sl, sr)) =
      nextr sr <&> \case
        Done () -> Skip $ Left (sl, Nothing)
        Skip s' -> Skip $ Right (Just sl, s')
        Yield (a, s') -> Yield (a, Left (sl, Just $ pure s'))
    next (Right (Nothing, sr)) =
      nextr sr <&> \case
        Done () -> Done ()
        Skip s' -> Skip $ Right (Nothing, s')
        Yield (a, s') -> Yield (a, Right (Nothing, s'))

  once (ListZerem z) = ListZerem $ Zerem.take 1 z

  lnot (ListZerem (Zerem next s)) = ListZerem $ Zerem next' (Just <$> s) where
    next' (Just s) =
      next s <&> \case
        Done () -> Yield ((), Nothing)
        Skip s' -> Skip (Just s')
        Yield _ -> Done ()
    next' Nothing = pure $ Done ()

  ifte (ListZerem (Zerem nexta msa)) k (ListZerem (Zerem nextb msb))
    = ListZerem $ Zerem next (IfteInit <$> msa) where
    next (IfteInit sa) =
      nexta sa >>= \case
        Done () -> msb <&> \sb -> Skip $ IfteFail sb
        Skip s' -> pure . Skip $ IfteInit s'
        Yield (a, s') -> case k a of -- RIP fusion
          ListZerem z -> Skip . IfteSucc s' <$> initZerem z
    next (IfteFail sb) =
      nextb sb <&> \case
        Done () -> Done ()
        Skip s' -> Skip $ IfteFail s'
        Yield (b, s') -> Yield (b, IfteFail s')
    next (IfteSucc sa (InitZerem nextb sb)) =
      nextb sb <&> \case
        Done () -> Skip $ IfteCont sa
        Skip s' -> Skip $ IfteSucc sa (InitZerem nextb s')
        Yield (b, s') -> Yield (b, IfteSucc sa (InitZerem nextb s'))
    next (IfteCont sa) =
      nexta sa >>= \case
        Done () -> pure $ Done ()
        Skip s' -> pure . Skip $ IfteCont s'
        Yield (a, s') -> case k a of
          ListZerem z -> Skip . IfteSucc s' <$> initZerem z

data ListZeremIfteStates sa sb zb a
  = IfteInit sa
  | IfteFail sb
  | IfteSucc sa zb
  | IfteCont sa

instance MonadTrans (ListZerem ()) where
  lift ma = ListZerem (Zerem (pure . next) (Just <$> ma)) where
    next Nothing = Done ()
    next (Just a) = Yield (a, Nothing)

instance MFunctor (ListZerem r) where
  hoist f (ListZerem (Zerem next s)) = ListZerem $ Zerem (f . next) (f s)

instance Comonad (ListZerem Void Identity) where
  extract (ListZerem (Zerem next ms)) = extract $ loop =<< ms where
    loop s =
      next s >>= \case
        Done v -> absurd v
        Skip s' -> loop s'
        Yield (a, _) -> pure a

  duplicate (ListZerem z@(Zerem next ms)) = ListZerem $ Zerem next' ((0 :: Int,) <$> ms) where
    next' (n, s) =
      next s <&> \case
        Done v -> absurd v
        Skip s' -> Skip (n, s')
        Yield (_, s') -> Yield (ListZerem $ Zerem.drop n z, (n+1, s'))

instance Foldable (ListZerem r Identity) where
  foldr f z (ListZerem stream) = runIdentity $ foldr_ f z stream
  {-# INLINE foldr #-}

  foldl' f z (ListZerem stream) = runIdentity $ foldl_ f z stream
  {-# INLINE foldl' #-}

instance Traversable (ListZerem () Identity) where
  traverse f = fmap (ListZerem . yieldMany) . traverse f . toList
  {-# INLINE traverse #-}

newtype Pipe u v f g m p a b = Pipe { runPipe :: p (Zerem (f a) m u) (Zerem (g b) m v) }

-- type Producer u v g m p b = forall f a. Pipe u v f g m p a b
-- type Consumer u v f m p a = forall g b. Pipe u v f g m p a b

instance Category p => Category (Pipe u u f f m p) where
  id = Pipe id
  (Pipe p2) . (Pipe p1) = Pipe $ p2 . p1

instance (Profunctor p, Bifunctor g, Functor m) => Functor (Pipe u v f g m p a) where
  fmap f = Pipe #. rmap (map f) .# runPipe

instance (Profunctor p, Bifunctor f, Bifunctor g, Functor m) => Profunctor (Pipe u v f g m p) where
  dimap lf rf = Pipe #. dimap (map lf) (map rf) .# runPipe
  lmap lf = Pipe #. lmap (map lf) .# runPipe
  rmap rf = Pipe #. rmap (map rf) .# runPipe

instance (Profunctor p, Bifunctor g, Functor m) => Choice (Pipe u v (,) g m p) where
  left' = Pipe #. dimap (mapMaybe eitherToMaybe) (map Left) .# runPipe where
    eitherToMaybe :: Either a c -> Maybe a
    eitherToMaybe = \case
      Left a -> Just a
      Right _c -> Nothing

  right' = Pipe #. dimap (mapMaybe eitherToMaybe) (map Right) .# runPipe where
    eitherToMaybe :: Either c a -> Maybe a
    eitherToMaybe = \case
      Left _c -> Nothing
      Right a -> Just a

instance ProfunctorFunctor (Pipe u v f g m) where
  promap f = Pipe #. f .# runPipe

instance ProfunctorMonad (Pipe Void Void (,) (,) Identity) where
  proreturn = Pipe #. dimap (extract .# ZipZerem) (runZipZerem #. pure)

  projoin =
    Pipe
    #. dimap
      (map $ runZipZerem #. pure)
      ((runZipZerem #. join)
      . (ZipZerem #. map ZipZerem))
    .# runPipe
    .# runPipe

map :: (Bifunctor f, Functor m) => (a -> b) -> Zerem (f a) m r -> Zerem (f b) m r
map f = maps (first f)
{-# INLINE map #-}

maps :: Functor m => (forall x. f x -> g x) -> Zerem f m r -> Zerem g m r
maps f (Zerem next s) = Zerem (fmap (mapYieldF f) . next) s
{-# INLINE maps #-}
-- {-# INLINE [0] maps #-}
-- {-# RULES "maps/coerce" [1] maps coerce = coerce #-}
-- {-# RULES "maps/maps"
--   forall
--     (phiA :: forall x. g x -> h x)
--     (phiB :: forall x. f x -> g x)
--     x.
--     maps phiA (maps phiB x) = maps (phiA . phiB) x #-}

mapM :: (Bitraversable f, Monad m) => (a -> m b) -> Zerem (f a) m r -> Zerem (f b) m r
mapM f = mapsM (bitraverse f pure)
{-# INLINE mapM #-}

mapsM :: Monad m => (forall x. f x -> m (g x)) -> Zerem f m r -> Zerem g m r
mapsM f (Zerem next s) = Zerem (mapYieldFM f <=< next) s
{-# INLINE mapsM #-}

flatLift :: Monad m => Zerem f m (m r) -> Zerem f m r
flatLift = flatMapLift id

flatMapLift :: Monad m => (r -> m r') -> Zerem f m r -> Zerem f m r'
flatMapLift k (Zerem next s) = Zerem next' s where
  next' s =
    next s >>= \case
      Done r -> Done <$> k r
      Skip s' -> pure $ Skip s'
      Yield fs' -> pure $ Yield fs'

(>>^) :: Monad m => Zerem f m r -> (r -> m r') -> Zerem f m r'
(>>^) = flip flatMapLift

hoistAny :: (forall x. m x -> n x) -> Zerem f m r -> Zerem f n r
hoistAny f (Zerem next s) = Zerem (f . next) (f s)

runZerem :: Monad m => Zerem m m r -> m r
runZerem (Zerem next s) = loop =<< s where
  loop s =
    next s >>= \case
      Done u    -> pure u
      Skip s'   -> loop s'
      Yield fs' -> fs' >>= loop

runZerem_ :: Monad m => Zerem m m r -> m ()
runZerem_ (Zerem next s) = loop =<< s where
  loop s =
    next s >>= \case
      Done _    -> pure ()
      Skip s'   -> loop s'
      Yield fs' -> fs' >>= loop

map_ :: (Comonad f, Monad m) => Zerem f m r -> m r
map_ = maps_ extract
{-# INLINE map_ #-}

maps_ :: Monad m => (forall x. f x -> x) -> Zerem f m r -> m r
maps_ phi (Zerem next s) = loop =<< s where
  loop s =
    next s >>= \case
      Done r -> pure r
      Skip s' -> loop s'
      Yield fs' -> loop (phi fs')
{-# INLINE maps_ #-}

mapM_ :: (Bitraversable f, Comonad (f b), Monad m) => (a -> m b) -> Zerem (f a) m r -> m r
mapM_ f = mapsM_ (fmap extract . bitraverse f pure)
{-# INLINE mapM_ #-}

mapsM_ :: Monad m => (forall x. f x -> m x) -> Zerem f m r -> m r
mapsM_ phi (Zerem next s) = loop =<< s where
  loop s =
    next s >>= \case
      Done r -> pure r
      Skip s' -> loop s'
      Yield fs' -> phi fs' >>= loop
{-# INLINE mapsM_ #-}

-- embed :: Functor m => (forall x. m x -> Zerem f n x) -> Zerem f m r -> Zerem f n r
-- embed f (Zerem next s done) = Zerem next' (s, Nothing) done where

foldr :: Monad m => (a -> b -> b) -> (r -> b) -> Zerem ((,) a) m r -> m b
foldr f !g (Zerem next s) = loop =<< s where
  loop s =
    next s >>= \case
      Done r -> pure (g r)
      Skip s' -> loop s'
      Yield (a, s') -> f a <$> loop s'

foldr_ :: Monad m => (a -> b -> b) -> b -> Zerem ((,) a) m r -> m b
foldr_ f z = Zerem.foldr f (const z)
{-# INLINE foldr_ #-}

foldrM :: Monad m => (a -> b -> m b) -> (r -> m b) -> Zerem ((,) a) m r -> m b
foldrM f !g (Zerem next s) = loop =<< s where
  loop s =
    next s >>= \case
      Done r -> g r
      Skip s' -> loop s'
      Yield (a, s') -> f a =<< loop s'

foldrM_ :: Monad m => (a -> b -> m b) -> m b -> Zerem ((,) a) m r -> m b
foldrM_ f z = Zerem.foldrM f (const z)
{-# INLINE foldrM_ #-}

foldl :: Monad m => (b -> a -> b) -> b -> (b -> r -> b') -> Zerem ((,) a) m r -> m b'
foldl f z !g (Zerem next s) = loop z =<< s where
  loop !z s =
    next s >>= \case
      Done r -> pure $! g z r
      Skip s' -> loop z s'
      Yield (a, s') -> loop (f z a) s'

foldl_ :: Monad m => (b -> a -> b) -> b -> Zerem ((,) a) m r -> m b
foldl_ f z = Zerem.foldl f z const
{-# INLINE foldl_ #-}

foldlM :: Monad m => (b -> a -> m b) -> m b -> (b -> r -> m b') -> Zerem ((,) a) m r -> m b'
foldlM f z !g (Zerem next s) = uncurry loop =<< liftA2 (,) z s where
  loop !z s =
    next s >>= \case
      Done r -> g z r
      Skip s' -> loop z s'
      Yield (a, s') -> (\z' -> loop z' s') =<< f z a

foldlM_ :: Monad m => (b -> a -> m b) -> m b -> Zerem ((,) a) m r -> m b
foldlM_ f z = Zerem.foldlM f z (const . pure)
{-# INLINE foldlM_ #-}

scanl :: (Functor m, Functor g) => (a -> x -> g x) -> x -> (r -> x -> x') -> Zerem ((,) a) m r -> Zerem g m x'
scanl f x g (Zerem next s) = Zerem next' ((,x) <$> s) where
  next' (s, !x) =
    next s <&> \case
      Done r -> Done (g r x)
      Skip s' -> Skip (s', x)
      Yield (a, s') -> Yield $ (s',) <$> f a x

scanl_ :: (Functor m, Functor g) => (a -> x -> g x) -> x -> Zerem ((,) a) m r -> Zerem g m x
scanl_ f x = Zerem.scanl f x (\_ x -> x)
{-# INLINE scanl_ #-}

scansl :: (Functor f, Functor m) => (forall s. f (s, x) -> g (s, x)) -> x -> (r -> x -> x') -> Zerem f m r -> Zerem g m x'
scansl f x g (Zerem next s) = Zerem next' ((,x) <$> s) where
  next' (s, !x) =
    next s <&> \case
      Done r -> Done (g r x)
      Skip s' -> Skip (s', x)
      Yield fs' -> Yield . f $ fmap (,x) fs'

scansl_ :: (Functor f, Functor m) => (forall s. f (s, x) -> g (s, x)) -> x -> Zerem f m r -> Zerem g m x
scansl_ f x = scansl f x (\_ x -> x)
{-# INLINE scansl_ #-}

scanMaybel :: (Functor m) => (a -> x -> (Maybe b, x)) -> x -> (r -> x -> x') -> Zerem ((,) a) m r -> Zerem ((,) b) m x'
scanMaybel f x g (Zerem next s) = Zerem next' ((,x) <$> s) where
  next' (s, !x) =
    next s <&> \case
      Done r -> Done (g r x)
      Skip s' -> Skip (s', x)
      Yield (a, s') -> case f a x of
        (Nothing, x') -> Skip  (s', x')
        (Just b , x') -> Yield (b, (s', x'))

scanMaybel_ :: (Functor m) => (a -> x -> (Maybe b, x)) -> x -> Zerem ((,) a) m r -> Zerem ((,) b) m x
scanMaybel_ f x = scanMaybel f x (\_ x -> x)
{-# INLINE scanMaybel_ #-}

catMaybes :: Functor m => Zerem ((,) (Maybe a)) m r -> Zerem ((,) a) m r
catMaybes (Zerem next s) = Zerem next' s where
  next' s =
    next s <&> \case
      Done r -> Done r
      Skip s' -> Skip s'
      Yield (Nothing, s') -> Skip s'
      Yield (Just a , s') -> Yield (a, s')
{-# INLINE catMaybes #-}

mapMaybe :: Functor m => (a -> Maybe b) -> Zerem ((,) a) m r -> Zerem ((,) b) m r
mapMaybe f (Zerem next s) = Zerem next' s where
  next' s =
    next s <&> \case
      Done r -> Done r
      Skip s' -> Skip s'
      Yield (a, s') -> case f a of
        Nothing -> Skip s'
        Just b  -> Yield (b, s')
{-# INLINE mapMaybe #-}

concat :: Applicative m => Zerem ((,) [a]) m r -> Zerem ((,) a) m r
concat (Zerem next ms) = Zerem next' (([],) <$> ms) where
  next' ([], s) =
    next s <&> \case
      Done r -> Done r
      Skip s' -> Skip ([], s')
      Yield (as, s') -> Skip (as, s')
  next' (a:as, s) = pure $ Yield (a, (as, s))
{-# INLINE concat #-}

concatMap :: Applicative m => (a -> [b]) -> Zerem ((,) a) m r -> Zerem ((,) b) m r
concatMap f (Zerem next ms) = Zerem next' (([],) <$> ms) where
  next' ([], s) =
    next s <&> \case
      Done r -> Done r
      Skip s' -> Skip ([], s')
      Yield (a, s') -> Skip (f a, s')
  next' (b:bs, s) = pure $ Yield (b, (bs, s))
{-# INLINE concatMap #-}

filter :: Functor m => (a -> Bool) -> Zerem ((,) a) m r -> Zerem ((,) a) m r
filter p (Zerem next s) = Zerem next' s where
  next' s =
    next s <&> \case
      Done r -> Done r
      Skip s' -> Skip s'
      Yield (a, s') -> if p a then Yield (a, s') else Skip s'
{-# INLINE filter #-}

take :: (Functor f, Applicative m) => Int -> Zerem f m r -> Zerem f m ()
take n (Zerem next s) = Zerem next' ((n,) <$> s) where
  next' (0, _) = pure (Done ())
  next' (n, s) =
    next s <&> \case
      Done _ -> Done ()
      Skip s' -> Skip (n, s')
      Yield fs' -> Yield $ (pred n,) <$> fs'

drop :: (Comonad f, Functor m) => Int -> Zerem f m r -> Zerem f m r
drop n z@(Zerem next s) = if n <= 0 then z else Zerem next' ((n,) <$> s) where
  next' (0, s) =
    next s <&> \case
      Done r -> Done r
      Skip s' -> Skip (0, s')
      Yield fs' -> Yield $ fmap (0,) fs'
  next' (n, s) =
    next s <&> \case
      Done r -> Done r
      Skip s' -> Skip (n, s')
      Yield fs' -> Skip $ (n-1, extract fs')

yield :: Applicative m => a -> Zerem ((,) a) m ()
yield = yieldMany . Identity
{-# INLINE yield #-}

yieldMany :: (Foldable f, Applicative m) => f a -> Zerem ((,) a) m ()
yieldMany = Zerem (pure . next) . pure . toList where
  next [] = Done ()
  next (a:as) = Yield (a, as)
{-# INLINE yieldMany #-}

yieldM :: Applicative m => m a -> Zerem ((,) a) m ()
yieldM = yieldManyM . Identity
{-# INLINE yieldM #-}

yieldManyM :: (Foldable f, Applicative m) => f (m a) -> Zerem ((,) a) m ()
yieldManyM = Zerem next . pure . toList where
  next [] = pure $ Done ()
  next (ma:as) = ma <&> Yield . (,as)
{-# INLINE yieldManyM #-}

concatStream
  :: (Functor m, Functor g)
  => (s -> m (Step s g u))
  -> Zerem ((,) s) m r
  -> Zerem g m r
concatStream inner (Zerem next s) = Zerem next' ((Nothing,) <$> s) where
  next' (Nothing, s) =
    next s <&> \case
      Done u -> Done u
      Skip s' -> Skip (Nothing, s')
      Yield (si, s') -> Skip (Just si, s')
  next' (Just si, s) =
    inner si <&> \case
      Done _ -> Skip (Nothing, s)
      Skip si' -> Skip (Just si', s)
      Yield gsi' -> Yield $ fmap ((,s) . Just) gsi'

mtimes
  :: (Monad m, Functor f)
  => Int -> Zerem f m r -> Zerem f m ()
mtimes n (Zerem next s0) =
  if n <= 0 then mempty
  else Zerem next' ((n,) <$> s0) where
  next' (n, s) =
    next s >>= \case
      -- Using Monad constraint to re-run the initialization action
      -- which might be creating a mutable reference that we don't want
      -- to re-use. Also, it's good to behave like (sconcat . replicate n).
      Done _ -> if n == 1 then pure $ Done () else Skip . (pred n,) <$> s0
      Skip s' -> pure $ Skip (n, s')
      Yield fs' -> pure . Yield $ fmap (n,) fs'