{-# LANGUAGE FunctionalDependencies, FlexibleInstances, MultiParamTypeClasses #-}
-- |A monad for consuming streams - I believe this is basically just a specialized version of the State monad. 
module Control.Monad.Consumer where

import Control.Applicative
import Control.Monad
import Control.Monad.Fix
import Control.Monad.Trans

-- * Consumer Monad

newtype Consumer c a = Consumer ([c] -> (a, [c]))

-- | Execute a stateful computation, as a result we get
-- the result of the computation, and the final state.
runConsumer :: [c] -> Consumer c a -> (a, [c])
runConsumer c (Consumer m) = m c

-- | Execute a stateful computation, ignoring the final state.
evalConsumer :: [c] -> Consumer c a -> a
evalConsumer c m     = fst (runConsumer c m)

-- | Execute a stateful computation, just for the side effect.
execConsumer        :: [c] -> Consumer c a -> [c]
execConsumer c m     = snd (runConsumer c m)

instance Functor (Consumer c) where
  fmap f (Consumer m) = Consumer (\c -> let (a,c1) = m c
                             in (f a, c1))

instance Monad (Consumer c) where
  return a         = Consumer (\c -> (a,c))
  Consumer m >>= k = Consumer (\c -> let (a,c1)     = m c
                                         Consumer n = k a
                                     in n c1)

instance MonadFix (Consumer c) where
  mfix f           = Consumer (\c -> 
                                   let Consumer m = f (fst r)
                                       r          = m c 
                                   in r)


class (Monad m) => MonadConsumer m c | m -> c where
    next :: (Monad m') => m (m' c) -- ^ return next element from stream
    peek :: (Monad m') => m (m' c) -- ^ peek at next element, but leave it in the stream
    poke :: c -> m c -- ^ push and element onto the beginning on the stream

instance MonadConsumer (Consumer c) c where
    next = Consumer $ \cs ->
           case cs of
             [] -> (fail "end of stream", [])
             (c:rest) -> (return c,rest)
    peek = Consumer $ \cs ->
           case cs of
             [] -> (fail "end of stream", [])
             cs@(c:_) -> (return c,cs)
    poke c = Consumer $ \cs -> (c,c:cs)

instance Applicative (Consumer a) where
    pure = return
    (<*>) = ap

-- * Consumer Monad Transformer

newtype ConsumerT c m a = ConsumerT { runConsumerT :: [c] -> m (a, [c]) }


instance (Monad m) => Functor (ConsumerT c m) where
	fmap f m = ConsumerT $ \c -> do
		(x, c') <- runConsumerT m c
		return (f x, c')

instance (Monad m) => Monad (ConsumerT c m) where
	return a = ConsumerT $ \c -> return (a, c)
	m >>= k  = ConsumerT $ \c -> do
		(a, c') <- runConsumerT m c
		runConsumerT (k a) c'
	fail str = ConsumerT $ \_ -> fail str

instance (MonadPlus m) => MonadPlus (ConsumerT c m) where
    mzero = lift mzero
    mplus m1 m2 = 
        ConsumerT $ \c ->
            let m1' = (runConsumerT m1) c
                m2' = (runConsumerT m2) c 
            in
              mplus m1' m2'

instance (Monad m) => MonadConsumer (ConsumerT c m) c where
    next = ConsumerT $ \cs ->
           case cs of
             [] -> return (fail "End of Stream", [])
             (c:rest) -> return (return c,rest)
    peek = ConsumerT $ \cs ->
           case cs of
             [] -> return (fail "End of Stream", [])
             cs@(c:_) -> return (return c,cs)
    poke c = ConsumerT $ \cs -> return (c,c:cs)

instance MonadTrans (ConsumerT c) where
    lift m = ConsumerT $ \c -> do
               a <- m
               return (a, c)

instance (MonadIO m) => MonadIO (ConsumerT c m) where
	liftIO = lift . liftIO

instance (Monad m) => Applicative (ConsumerT c m) where
    pure = return
    (<*>) = ap