{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE DeriveGeneric       #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE EmptyCase           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}
{-# LANGUAGE TemplateHaskellQuotes #-}
module Staged.Stream.Combinators (
    -- * Construction
    singleton,
    fromList,
    fromListM,
    unfold,
    iterate,
    iterateM,
    replicate,
    replicateM,
    -- * Transformations
    map,
    mapWithInput,
    lmap,
    traverse,
    filter,
    filterM,
    drop,
    -- * Append
    append,
    empty,
    -- * Zipping
    zipWith,
    zipWithM,
    repeat,
    repeatM,
    -- * Aligning
    -- * Recursion
    bfsTreeM,
    -- * Pipes
    idPipe,
    mapPipe,
    traversePipe,
    filterPipe,
    -- * Elimination
    run,
    foldl,
    foldlM,
    toList,
    -- * Conversion
    fromPure,
    toPure,
    ) where

import Prelude (($), Bool(..), Either (..), Int, subtract, (<), (+), (.), Maybe (..), Monad, id)
-- import Prelude (undefined)

import Generics.SOP (SOP (..), SListI2)
import Data.Functor.Identity (Identity (..))

import qualified Control.Category as  C

import Staged.Commons
import Staged.Stream.Step
import Staged.Stream.States
import Staged.Stream.Pure.Type
import Staged.Stream.Type
import Staged.Stream.Convenience

import qualified Staged.Stream.Pure as Pure

-------------------------------------------------------------------------------
-- Conversion
-------------------------------------------------------------------------------

-- | Convert from pure 'Stream' to monadic 'StreamM'.
fromPure :: Stream a b -> StreamM m a b
fromPure (MkStream start step) = MkStreamM start step

-- | Convert from @'StreamM' 'Identity'@ to 'Stream'.
--
-- /Note:/ while 'fromPure' conversion is /clean/,
-- 'toPure' will add 'runIdentity' and 'Identity' into the splices.
--
toPure :: StreamM Identity a b -> Stream a b
toPure (MkStreamM start steps) = MkStream start $ \curr k ->
    [|| runIdentity ||] @@ (steps curr $ \s -> C [|| Identity ||] @@ k s )

-------------------------------------------------------------------------------
-- Construction
-------------------------------------------------------------------------------

-- |
--
-- @
-- 'singleton' :: C b -> 'StreamM' m a b
-- 'singleton' b = 'map' ([|| const ||] @@ b)
-- @
singleton :: forall b a m cb. IsCode b cb => cb -> StreamM m a b
singleton = fromPure . Pure.singleton

-- |
--
-- @
-- 'fromList' :: (C a -> C [b]) -> 'StreamM' m a b
-- @
fromList :: ToCodeFn a [b] fn => fn -> StreamM m a b
fromList = fromPure . Pure.fromList

-- |
--
-- @
-- 'fromListM' :: (C a -> C (m [b])) -> 'StreamM' m a b
-- @
--
fromListM :: forall a b m fn. (Monad m, ToCodeFn a (m [b]) fn) => fn -> StreamM m a b
fromListM f = mkStreamM (Left . toFn f) steps where
    steps
        :: Either (C (m [b])) (C [b])
        -> (Step (C b) (Either (C (m [b])) (C [b])) -> C (m r))
        -> C (m r)
    steps (Left mbs) k = mbs >>>= \bs -> k (Skip (Right bs))
    steps (Right mb) k = scaseList mb
        (k Stop)
        (\b bs' -> k (Emit b (Right bs')))

-- |
--
-- Note: this can be used to unfold with monadic effects too.
--
-- @
-- 'unfold' :: (C a -> C b) -> (C b -> CPS (Maybe (C c, C b))) -> 'StreamM' m a c
-- @
unfold
    :: forall a b c m fn. ToCodeFn a b fn
    => fn
    -> (forall r. C b -> (Maybe (C c, C b) -> C (m r)) -> C (m r)) -- ^ unfolding
    -> StreamM m a c
unfold start f = mkStreamM (toFn start) steps where
    steps :: C b -> (Step (C c) (C b) -> C (m r)) -> C (m r)
    steps curr k = f curr $ \case
        Nothing        -> k Stop
        Just (c, next) -> k (Emit c next)

-- |
--
-- @
-- 'iterate' :: (C a -> C a) -> 'StreamM' m a a
-- @
iterate :: (ToCodeFn a a endo) => endo -> StreamM m a a
iterate f = fromPure (Pure.iterate f)

-- |
--
-- @
-- 'iterateM' :: (C a -> C (m a)) -> C a -> 'StreamM' m i a
-- @
iterateM :: (Monad m, ToCodeFn a (m a) endo) => endo -> StreamM m a a
iterateM f = mkStreamM id $ \curr k ->
    toFn f curr >>>= \x ->
    k (Emit curr x)

-- |
--
-- @
-- 'replicate' :: C Int -> C a -> 'Stream' i a
-- @
replicate :: (IsCode a ca, IsCode Int ci) => ci -> ca -> StreamM m i a
replicate i a = take i (repeat a)

-- |
--
-- @
-- 'replicate' :: C Int -> C a -> 'Stream' i a
-- @
replicateM :: (Monad m, IsCode (m a) ca, IsCode Int ci) => ci -> ca -> StreamM m i a
replicateM i a = take i (repeatM a)

-------------------------------------------------------------------------------
-- Transformations
-------------------------------------------------------------------------------

-- |
-- 
-- @
-- 'map' :: (C b -> C c) -> 'Stream' a b -> 'Stream' a c
-- @
map :: forall a b c m fn. ToCodeFn a b fn
    => fn -> StreamM m c a -> StreamM m c b
map f (MkStreamM s0 steps0) = MkStreamM s0 (go steps0) where
    go :: (SOP C xss -> (Step (C a) (SOP C xss) -> C r) -> C r)
       -> (SOP C xss -> (Step (C b) (SOP C xss) -> C r) -> C r)
    go steps s k = steps s $ \case
        Stop      -> k Stop
        Skip   s' -> k (Skip s')
        Emit a s' -> k (Emit (toFn f a) s')

-- |
--
-- @
-- 'mapWithInput' :: (C a -> C b -> C c) -> 'StreamM' m a b -> 'StreamM' m a c
-- @
mapWithInput
    :: forall a b c m fn. ToCodeFn2 a b c fn
    => fn -> StreamM m a b -> StreamM m a c
mapWithInput f (MkStreamM s0 steps0) =
    mkStreamM (\a -> (a, s0 a)) $ \(a, curr) k -> steps0 curr $ \case
        Stop        -> k Stop
        Skip   next -> k (Skip (a, next))
        Emit b next -> k (Emit (toFn2 f a b) (a, next))

-- |
--
-- @
-- 'lmap' :: (C a -> C b) -> 'StreamM' m b c -> 'StreamM' m a c
-- @
lmap
    :: forall a b c m fn. ToCodeFn a b fn
    => fn -> StreamM m b c -> StreamM m a c
lmap f (MkStreamM s0 steps0) = MkStreamM (s0 . toFn f) steps0

-- |
--
-- @
-- 'traverse' :: (C b -> C (m c)) -> 'StreamM' m a b -> 'StreamM' m a c
-- @
traverse
    :: forall a b c m fn. (Monad m, ToCodeFn b (m c) fn)
    => fn -> StreamM m a b -> StreamM m a c
traverse f (MkStreamM s0 steps0) = MkStreamM s0 (go steps0) where
    go :: (SOP C xss -> (Step (C b) (SOP C xss) -> C (m r)) -> C (m r))
       -> (SOP C xss -> (Step (C c) (SOP C xss) -> C (m r)) -> C (m r))
    go steps s k = steps s $ \case
        Stop      -> k Stop
        Skip   s' -> k (Skip s')
        Emit a s' -> toFn f a >>>= \a' -> k (Emit a' s')

-- TODO: traverseWithInput
-- TODO: ltraverse

-- |
--
-- @
-- 'filter' :: (C a -> C Bool) -> 'StreamM' m c a -> 'StreamM' m c a
-- @
filter
    :: forall a c m predicate. (ToCodeFn a Bool predicate)
    => predicate -> StreamM m c a -> StreamM m c a
filter p (MkStreamM s0 steps0) = MkStreamM s0 (go steps0) where
    go :: (SOP C xss -> (Step (C a) (SOP C xss) -> C r) -> C r)
       -> (SOP C xss -> (Step (C a) (SOP C xss) -> C r) -> C r)
    go steps s k = steps s $ \case
        Stop      -> k Stop
        Skip   s' -> k (Skip s')
        Emit a s' -> sIfThenElse
            (toFn p a)
            (k (Emit a s'))
            (k (Skip s'))

-- |
--
-- @
-- 'filterM' :: (C a -> C (m Bool)) -> 'StreamM' m c a -> 'StreamM' m c a
-- @
filterM
    :: forall a c m predicateM. (Monad m, ToCodeFn a (m Bool) predicateM)
    => predicateM -> StreamM m c a -> StreamM m c a
filterM p (MkStreamM s0 steps0) = MkStreamM s0 (go steps0) where
    go :: (SOP C xss -> (Step (C a) (SOP C xss) -> C (m r)) -> C (m r))
       -> (SOP C xss -> (Step (C a) (SOP C xss) -> C (m r)) -> C (m r))
    go steps s k = steps s $ \case
        Stop      -> k Stop
        Skip   s' -> k (Skip s')
        Emit a s' -> toFn p a >>>= \pa -> sIfThenElse
            pa
            (k (Emit a s'))
            (k (Skip s'))

-------------------------------------------------------------------------------
-- Take and drop
-------------------------------------------------------------------------------

-- |
--
-- @
-- 'take' :: C Int -> 'StreamM' m a b -> 'StreamM' m a b
-- @
take :: IsCode Int n => n -> StreamM m a b -> StreamM m a b
take n (MkStreamM start steps) = do
    mkStreamM (\a -> (C [|| 0 ||], start a)) $ \(i, xss) k -> steps xss $ \case
        Stop        -> k Stop
        Skip   xss' -> k (Skip (i, xss'))
        Emit b xss' -> sIfThenElse
            (C [|| (<) ||] @@ i @@ toCode n)
            (k (Emit b (C [|| (1 +) ||] @@ i, xss')))
            (k Stop)
-- |
--
-- @
-- 'drop' :: C Int -> 'StreamM' m a b -> 'StreamM' m a b
-- @
drop :: IsCode Int n => n -> StreamM m a b -> StreamM m a b
drop n (MkStreamM start steps) =
    mkStreamM (\a -> DropL (toCode n) (start a)) $ \step k -> case step of
        DropL m xss -> steps xss $ \case
            Stop        -> k Stop
            Skip   xss' -> k (Skip (DropL m xss'))
            Emit b xss' -> sIfThenElse
                (C [|| (0 <) ||] @@ m)
                (k (Skip   (DropL (C [|| subtract 1 ||] @@ m) xss')))
                (k (Emit b (DropR xss')))

        DropR xss -> steps xss $ \case
            Stop        -> k Stop
            Skip   xss' -> k (Skip   (DropR xss'))
            Emit b xss' -> k (Emit b (DropR xss'))

-------------------------------------------------------------------------------
-- Append
-------------------------------------------------------------------------------

append :: forall a b m. StreamM m a b -> StreamM m a b -> StreamM m a b
append (MkStreamM startL stepsL) (MkStreamM startR stepsR) =
    mkStreamM (\a -> AppL a (startL a)) $ \step k -> case step of
        AppL a xss -> stepsL xss $ \case
            Stop        -> k (Skip   (AppR (startR a)))
            Skip   xss' -> k (Skip   (AppL a xss'))
            Emit b xss' -> k (Emit b (AppL a xss'))

        AppR yss -> stepsR yss $ \case
            Stop -> k Stop
            Skip   yss' -> k (Skip   (AppR yss'))
            Emit b yss' -> k (Emit b (AppR yss'))
        
empty :: StreamM m a b
empty = fromPure Pure.empty

-------------------------------------------------------------------------------
-- Zipping
-------------------------------------------------------------------------------

-- |
--
-- @
-- 'zipWith' :: (C a -> C b -> C c) -> 'StreamM' m i a -> 'StreamM' m i b -> 'StreamM' m i c
-- @
zipWith :: forall i a b c m abc. ToCodeFn2 a b c abc => abc -> StreamM m i a -> StreamM m i b -> StreamM m i c
zipWith h (MkStreamM start0 steps0) (MkStreamM start1 steps1) =
    mkStreamM (\i -> ZipL (start0 i) (start1 i)) (steps steps0 steps1)
  where
    steps
        :: (forall r'. SOP C xss -> (Step (C a) (SOP C xss) -> C (m r')) -> C (m r'))
        -> (forall r'. SOP C yss -> (Step (C b) (SOP C yss) -> C (m r')) -> C (m r'))
        -> Zip a xss yss
        -> (Step (C c) (Zip a xss yss) -> C (m r))
        -> C (m r)
    steps f _ (ZipL xss yss) k = f xss $ \case
        Stop        -> k Stop
        Skip   xss' -> k (Skip (ZipL xss' yss))
        Emit a xss' -> k (Skip (ZipR a xss' yss))

    steps _ g (ZipR a xss yss) k = g yss $ \case
        Stop        -> k Stop
        Skip   yss' -> k (Skip               (ZipR a xss yss'))
        Emit b yss' -> k (Emit (toFn2 h a b) (ZipL xss yss'))

-- |
--
-- @
-- 'zipWithM' :: (C a -> C b -> C (m c)) -> 'StreamM' m i a -> 'StreamM' m i b -> 'StreamM' m i c
-- @
zipWithM
    :: forall i a b c m abc. (Monad m, ToCodeFn2 a b (m c) abc)
    => abc -> StreamM m i a -> StreamM m i b -> StreamM m i c
zipWithM h (MkStreamM start0 steps0) (MkStreamM start1 steps1) =
    mkStreamM (\i -> ZipL (start0 i) (start1 i)) (steps steps0 steps1)
  where
    steps
        :: (forall r'. SOP C xss -> (Step (C a) (SOP C xss) -> C (m r')) -> C (m r'))
        -> (forall r'. SOP C yss -> (Step (C b) (SOP C yss) -> C (m r')) -> C (m r'))
        -> Zip a xss yss
        -> (Step (C c) (Zip a xss yss) -> C (m r))
        -> C (m r)
    steps f _ (ZipL xss yss) k = f xss $ \case
        Stop        -> k Stop
        Skip   xss' -> k (Skip (ZipL xss' yss))
        Emit a xss' -> k (Skip (ZipR a xss' yss))

    steps _ g (ZipR a xss yss) k = g yss $ \case
        Stop        -> k Stop
        Skip   yss' -> k (Skip (ZipR a xss yss'))
        Emit b yss' -> toFn2 h a b >>>= \c -> k (Emit c (ZipL xss yss'))


-- |
--
-- @
-- 'repeat' :: 'C' a -> 'StreamM' m i a
-- @
repeat :: IsCode a ca => ca -> StreamM m i a
repeat = fromPure . Pure.repeat

-- |
--
-- @
-- 'repeatM' :: 'C' (m a) -> 'StreamM' m i a
-- @
repeatM :: (Monad m, IsCode (m a) cma) => cma -> StreamM m i a
repeatM a = mkStreamM (\_ -> ()) $ \() k -> toCode a >>>= \a' -> k (Emit a' ())

-------------------------------------------------------------------------------
-- Recursion
-------------------------------------------------------------------------------

-- |
-- 
-- @
-- 'bfsTreeM' :: 'StreamM' m a a -> (C a -> C (m Bool)) -> 'StreamM' m a a
-- @
bfsTreeM
    :: forall a m predicateM. (Monad m, ToCodeFn a (m Bool) predicateM)
    => StreamM m a a  -- ^ endo-stream
    -> predicateM     -- ^ whether to recurse on a produced element
    -> StreamM m a a
bfsTreeM (MkStreamM start0 steps0) p = mk start0 steps0 where
    mk  :: forall xss. SListI2 xss
        => (C a -> SOP C xss)
        -> (forall r. SOP C xss -> (Step (C a) (SOP C xss) -> C (m r)) -> C (m r))
        -> StreamM m a a
    mk start1 steps1 = mkStreamM (BfsNext . start1) steps2 where
        steps2  :: BFS a xss -> (Step (C a) (BFS a xss) -> C (m r)) -> C (m r)
        steps2 (BfsNext curr) k = steps1 curr $ \case
            Stop        -> k Stop
            Skip   next -> k (Skip (BfsNext next))
            Emit a next -> toFn p a >>>= \b -> sIfThenElse b
                (k (Emit a (BfsStep (scons a snil) next)))
                (k (Emit a (BfsNext next)))

        steps2 (BfsStep as curr) k = steps1 curr $ \case
            Stop -> scaseList as
                (k Stop)
                (\a as' -> k (Skip (BfsStep as' (start1 a))))
            Skip   next -> k (Skip (BfsStep as next))
            Emit a next -> toFn p a >>>= \b -> sIfThenElse b
                (k (Emit a (BfsStep (scons a as) next)))
                (k (Emit a (BfsStep as next)))

-------------------------------------------------------------------------------
-- Pipes
-------------------------------------------------------------------------------

-- | Identity stream.
idPipe :: StreamM m a a
idPipe = C.id

-- | Similar to 'map', prefer using 'map'.
--
-- @
-- 'mapPipe' :: (C a -> C b) -> 'StreamM' m a b
-- @
mapPipe :: forall a b m fn. ToCodeFn a b fn => fn -> StreamM m a b
mapPipe f = mkStreamM (Just . toFn f) step where
    step :: Maybe (C b) -> (Step (C b) (Maybe (C b)) -> C r) -> C r
    step Nothing  k = k Stop
    step (Just b) k = k (Emit b Nothing)

-- | Similar to 'traverse', prefer using 'traverse'.
--
-- @
-- 'traversePipe' :: (C a -> C (m b)) -> 'StreamM' m a b
-- @
traversePipe :: forall a b m fn. (Monad m, ToCodeFn a (m b) fn) => fn -> StreamM m a b
traversePipe f = mkStreamM (Just . toFn f) step where
    step :: Maybe (C (m b)) -> (Step (C b) (Maybe (C (m b))) -> C (m r)) -> C (m r)
    step Nothing   k = k Stop
    step (Just mb) k = mb >>>= \b -> k (Emit b Nothing)

-- | Similar to 'filter', prefer using 'filter'.
-- 
-- @
-- 'filterPipe' :: (C a -> C Bool) -> 'StreamM' m a a
-- @
filterPipe :: forall a m predicate. ToCodeFn a Bool predicate => predicate -> StreamM m a a
filterPipe p = mkStreamM Just step where
    step :: Maybe (C a) -> (Step (C a) (Maybe (C a)) -> C r) -> C r
    step Nothing  k = k Stop
    step (Just a) k = sIfThenElse
        (toFn p a)
        (k (Emit a Nothing))
        (k Stop) -- stop immediately, don't go to the second state.

-------------------------------------------------------------------------------
-- Elimination
-------------------------------------------------------------------------------

-- |
--
-- @
-- 'run' :: C a -> 'StreamM' m a b -> GHCCode (m ())
-- @
run :: forall a b m ca. (Monad m, IsCode a ca)
    => ca -> StreamM m a b -> GHCCode (m ())
run z (MkStreamM start steps0) =
    fromCode $ sletrec_SOP (body steps0) (start (toCode z))
  where
    body
        :: (SOP C xss -> (Step (C b) (SOP C xss) -> C (m ())) -> C (m ()))
        -> (SOP C xss -> C (m ()))
        -> SOP C xss -> C (m ())
    body steps loop curr = steps curr $ \case
        Stop        -> sreturn sunit
        Skip   next -> loop next
        Emit _ next -> loop next

-- |
--
-- @
-- 'foldl' :: (C r -> C b -> C r) -> C r -> C a -> StreamM m a b -> GHCCode (m r)
-- @
foldl :: forall r a b m fn init start. (Monad m, IsCode r init, IsCode a start, ToCodeFn2 r b r fn)
      => fn -> init -> start -> StreamM m a b -> GHCCode (m r)
foldl op e z (MkStreamM xs steps0) =
    fromCode $ sletrec1_SOP (body steps0) (xs (toCode z)) (toCode e)
  where
    body
        :: (SOP C xss -> (Step (C b) (SOP C xss) -> C (m r)) -> C (m r))
        -> (SOP C xss -> C r -> C (m r))
        -> (SOP C xss -> C r -> C (m r))
    body steps loop curr acc = steps curr $ \case
        Stop        -> sreturn acc
        Skip   next -> loop next acc
        Emit b next -> loop next (toFn2 op acc b)

-- |
--
-- @
-- 'foldlM' :: (C r -> C b -> C (m r)) -> C r -> C a -> StreamM m a b -> GHCCode (m r)
-- @
foldlM
    :: forall r a b m fn init start. (Monad m, IsCode r init, IsCode a start, ToCodeFn2 r b (m r) fn)
    => fn -> init -> start -> StreamM m a b -> GHCCode (m r)
foldlM op e z (MkStreamM xs steps0) =
    fromCode $ sletrec1_SOP (body steps0) (xs (toCode z)) (toCode e)
  where
    body
        :: (SOP C xss -> (Step (C b) (SOP C xss) -> C (m r)) -> C (m r))
        -> (SOP C xss -> C r -> C (m r))
        -> (SOP C xss -> C r -> C (m r))
    body steps loop curr acc = steps curr $ \case
        Stop        -> sreturn acc
        Skip   next -> loop next acc
        Emit b next -> toFn2 op acc b >>>= loop next
 
-- |
--
-- @
-- 'toList' :: C a -> StreamM a b -> GHCCode (m [b])
-- @
toList
    :: forall a b m ca. (Monad m, IsCode a ca)
    => ca -> StreamM m a b -> GHCCode (m [b])
toList a (MkStreamM start steps0) =
    fromCode $ sletrec_SOP (body steps0) (start (toCode a))
  where
    body
        :: (SOP C xss -> (Step (C b) (SOP C xss) -> C (m [b])) -> C (m [b]))
        -> (SOP C xss -> C (m [b]))
        -> SOP C xss -> C (m [b])
    body steps loop curr = steps curr $ \case
        Stop        -> sreturn snil
        Skip   next -> loop next
        Emit b next -> sfmap (scons b) (loop next)