{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE NoImplicitPrelude #-}

module Streaming.NonEmpty where

import Control.Category (id, (.))
import Control.Monad (Monad (return))
import Data.Bool (Bool)
import Data.Either
import Data.Eq ((==))
import Data.Function (($))
import Data.Functor (Functor, fmap, (<$>))
import Data.List.NonEmpty hiding (groupBy, map)
import Data.Semigroup (Semigroup, (<>))
import qualified Data.Semigroup as Semigroup
import Streaming (Of ((:>)), Stream, first, runIdentity, wrap)
import Streaming.Internal
import qualified Streaming.Prelude as S
import Prelude hiding ((.), map, id)

-- | A stream with at least one functorial layer
newtype NEStream f m a = NEStream (f (Stream f m a)) deriving (Functor)

-- | Is a stream
toStream :: (Functor f, Monad m) => NEStream f m r -> Stream f m r
toStream (NEStream f) = wrap f

map :: Monad m => (t -> b) -> NEStream (Of t) m a -> NEStream (Of b) m a
map f (NEStream (x :> s)) = NEStream (f x :> S.map f s)


-- | fold1 is safe
fold1 :: Monad m => (a -> a -> a) -> NEStream (Of a) m r -> m (Of a r)
fold1 f (NEStream (x :> s)) = S.fold f x id s

-- | sconcat as in 'Data.Semigroup', with result
sconcat :: (Semigroup a, Monad m) => NEStream (Of a) m r -> m (Of a r)
sconcat = fold1 (<>)

-- | harvest the list with the result
toNonEmpty :: Monad m => NEStream (Of a) m r -> m (Of (NonEmpty a) r)
toNonEmpty (NEStream (x :> r)) = first (x :|) <$> S.toList r

-- | harvest the list only
toNonEmpty_ :: Monad m => NEStream (Of a) m r -> m (NonEmpty a)
toNonEmpty_ = fmap S.fst' . toNonEmpty

-- | sconcat as in 'Data.Semigroup'
sconcat_ :: (Monad m, Semigroup a) => NEStream (Of a) m r -> m a
sconcat_ (NEStream (x :> s)) = S.fold_ (<>) x id s

-- | group by some  equality  in non empty groups
groupBy
  :: (Monad m)
  => (a -> a -> Bool)
  -> Stream (Of a) m r
  -> Stream (NEStream (Of a) m) m r
groupBy equals = loop
  where
    loop stream = Effect $ do
      e <- S.next stream
      return $ case e of
        Left r -> Return r
        Right (a, p') ->
          Step $
            fmap loop (NEStream (a :> S.span (equals a) p'))

-- | what groupBy could be
groupByPure :: (a -> a -> Bool) -> [a] -> [NonEmpty a]
groupByPure equals = runIdentity . S.toList_ . S.mapped toNonEmpty . groupBy equals . S.each

-- | collapse semigroups by some equality
groupSemigroupBy :: (Semigroup a, Monad m) => (a -> a -> Bool) -> Stream (Of a) m r -> Stream (Of a) m r
groupSemigroupBy f = S.mapped sconcat . groupBy f

-- | what should be possible to do with 'groupBy'
-- fmap sconcat . groupBy equals
groupSemigroupByPure :: Semigroup b => (b -> b -> Bool) -> [b] -> [b]
groupSemigroupByPure equals = runIdentity . S.toList_ . groupSemigroupBy equals . S.each