{-# LANGUAGE BangPatterns, ExistentialQuantification #-}
module Data.RangeMin.Fusion.Stream.Monadic where

import Data.Vector.Fusion.Util
import Control.Monad hiding (replicateM)
import Prelude hiding (replicate, map)

data Stream m a = forall s . Stream (Int -> s -> m (Yield s a)) s {-# UNPACK #-} !Int
data Yield s a = Done | Last a | Yield a s

{-# INLINE length #-}
length :: Stream m a -> Int
length (Stream _ _ n) = n

{-# INLINE iunfoldMN #-}
iunfoldMN :: Monad m => Int -> (Int -> b -> m (Maybe (a, b))) -> b -> Stream m a
iunfoldMN !n f s0 = Stream suc s0 n where
	suc i s = do
	  unf <- f i s
	  case unf of
	    Nothing	 -> return Done
	    Just (a, s') -> return (Yield a s')

{-# INLINE unfoldMN #-}
unfoldMN :: Monad m => Int -> (b -> m (Maybe (a, b))) -> b -> Stream m a
unfoldMN n = iunfoldMN n . const

{-# INLINE iunfoldN #-}
iunfoldN :: Monad m => Int -> (Int -> b -> Maybe (a, b)) -> b -> Stream m a
iunfoldN n f = iunfoldMN n (\ i b -> return (f i b))

{-# INLINE unfoldN #-}
unfoldN :: Monad m => Int -> (b -> Maybe (a, b)) -> b -> Stream m a
unfoldN n = iunfoldN n . const

{-# INLINE imapM #-}
imapM :: Monad m => (Int -> a -> m b) -> Stream m a -> Stream m b
imapM f (Stream suc s0 n) = Stream suc' s0 n where
	suc' i s = do
	  step <- suc i s
	  case step of
	    Done	-> return Done
	    Last a	-> Last `liftM` f i a
	    Yield a s'	-> do
		    b <- f i a
		    return (Yield b s')

{-# INLINE imap #-}
imap :: Monad m => (Int -> a -> b) -> Stream m a -> Stream m b
imap f = imapM (\ i a -> return (f i a))

{-# INLINE map #-}
map :: Monad m => (a -> b) -> Stream m a -> Stream m b
map = imap . const

{-# INLINE replicateM #-}
replicateM :: Monad m => Int -> m a -> Stream m a
replicateM !n m = Stream (\ i _ -> 
	if i < n then do
		a <- m
		return (Yield a ())
	  else return Done) () n

{-# INLINE replicate #-}
replicate :: Monad m => Int -> a -> Stream m a
replicate n a = replicateM n (return a)

{-# INLINE enumN #-}
enumN :: Monad m => Int -> Stream m Int
enumN n = imap const (replicate n ())

{-# INLINE enumNR #-}
enumNR :: Monad m => Int -> Stream m Int
enumNR !n = Stream suc n n where
	suc _ 0 = return Done
	suc _ i = let !i' = i - 1 in return (Yield i' i')

{-# INLINE cap #-}
cap :: Monad m => Stream m a -> Stream m a
cap (Stream suc s0 n) = Stream (\ i s -> if i < n then suc i s else return Done) s0 n

{-# INLINE generateM #-}
generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
generateM n f = imapM (const . f) (replicate n ())

{-# INLINE generate #-}
generate :: Monad m => Int -> (Int -> a) -> Stream m a
generate n f = generateM n (return . f)

{-# INLINE imapM_ #-}
imapM_ :: Monad m => (Int -> a -> m b) -> Stream m a -> m ()
imapM_ f = consume . imapM f

{-# INLINE mapM_ #-}
mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
mapM_ = imapM_ . const

{-# INLINE consume #-}
consume :: Monad m => Stream m a -> m ()
consume (Stream suc s0 _) = go 0 s0 where
	go !i s = do
	  step <- suc i s
	  case step of
	    Yield _ s' -> go (i+1) s'
	    _ -> return ()

{-# INLINE snocM #-}
snocM :: Monad m => Stream m a -> m a -> Stream m a
Stream suc s0 n `snocM` z = Stream suc' s0 (n+1) where
	suc' i s = if i < n then suc i s else Last `liftM` z

{-# INLINE snoc #-}
snoc :: Monad m => Stream m a -> a -> Stream m a
stream `snoc` z = stream `snocM` return z

{-# INLINE izipWithM #-}
izipWithM :: Monad m => (Int -> a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
izipWithM f (Stream suc1 s01 n1) (Stream suc2 s02 n2) = Stream suc (s01, s02) (min n1 n2) where
	suc i (s1, s2) = do
		!step1 <- suc1 i s1
		!step2 <- suc2 i s2
		case (step1, step2) of
			(Last a, Last b) -> Last `liftM` f i a b
			(Last a, Yield b _) -> Last `liftM` f i a b
			(Yield a _, Last b) -> Last `liftM` f i a b
			(Yield a s1', Yield b s2') -> do
				c <- f i a b
				return (Yield c (s1', s2'))
			(Done, _)	-> return Done
			(_, Done)	-> return Done


{-# INLINE izipWith #-}
izipWith :: Monad m => (Int -> a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
izipWith f = izipWithM (\ i a b -> return (f i a b))

{-# INLINE zipWithM #-}
zipWithM :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
zipWithM = izipWithM . const

{-# INLINE zipWith #-}
zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
zipWith = izipWith . const

{-# INLINE ifoldlM #-}
ifoldlM :: Monad m => (b -> Int -> a -> m b) -> b -> Stream m a -> m b
ifoldlM f z0 (Stream suc s0 _) = go 0 s0 z0 where
	go !i s z = do
	  step <- suc i s
	  case step of
	    Done -> return z
	    Last a -> f z i a
	    Yield a s' ->
	    	go (i+1) s' =<< f z i a

{-# INLINE foldlM #-}
foldlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> m b
foldlM f = ifoldlM (\ z _ a -> f z a)

{-# INLINE ifoldl #-}
ifoldl :: Monad m => (b -> Int -> a -> b) -> b -> Stream m a -> m b
ifoldl f = ifoldlM (\ z i a -> return (f z i a))

{-# INLINE foldl #-}
foldl :: Monad m => (b -> a -> b) -> b -> Stream m a -> m b
foldl f = ifoldl (\ z _ a -> f z a)

{-# INLINE imapAccumLM #-}
imapAccumLM :: Monad m => (b -> Int -> a -> m (c, b)) -> b -> Stream m a -> Stream m c
imapAccumLM f z0 (Stream suc s0 n) = Stream suc' (s0, z0) n where
	suc' i (s, z) = do
	  step <- suc i s
	  case step of
	    Done -> return Done
	    Last a -> (Last . fst) `liftM` f z i a
	    Yield a s' -> do
	    	(y, z') <- f z i a
	    	return (Yield y (s', z'))

{-# INLINE mapAccumLM #-}
mapAccumLM :: Monad m => (b -> a -> m (c, b)) -> b -> Stream m a -> Stream m c
mapAccumLM f = imapAccumLM (\ z _ a -> f z a)

{-# INLINE imapAccumL #-}
imapAccumL :: Monad m => (b -> Int -> a -> (c, b)) -> b -> Stream m a -> Stream m c
imapAccumL f = imapAccumLM (\ z i a -> return (f z i a))

{-# INLINE mapAccumL #-}
mapAccumL :: Monad m => (b -> a -> (c, b)) -> b -> Stream m a -> Stream m c
mapAccumL f = imapAccumL (\ z _ a -> f z a)

{-# INLINE ipostscanlM #-}
ipostscanlM :: Monad m => (b -> Int -> a -> m b) -> b -> Stream m a -> Stream m b
ipostscanlM f = imapAccumLM (\ z i a -> liftM (\ a -> (a, a)) (f z i a))

{-# INLINE ipostscanl #-}
ipostscanl :: Monad m => (b -> Int -> a -> b) -> b -> Stream m a -> Stream m b
ipostscanl f = ipostscanlM (\ z i a -> return (f z i a))

{-# INLINE postscanlM #-}
postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
postscanlM f = ipostscanlM (\ z _ a -> f z a)

{-# INLINE postscanl #-}
postscanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
postscanl f = ipostscanl (\ z _ a -> f z a)

{-# INLINE iterateMN #-}
iterateMN :: Monad m => Int -> a -> (a -> m a) -> Stream m a
iterateMN !n z f = cap $ Stream suc z n where
	suc _ z = do
		z' <- f z
		return (Yield z' z')

{-# INLINE iterateN #-}
iterateN :: Monad m => Int -> a -> (a -> a) -> Stream m a
iterateN n z f = iterateMN n z (return . f)

{-# INLINE unbox #-}
unbox :: Monad m => Stream m (Box a) -> Stream m a
unbox (Stream suc s0 n) = Stream suc' s0 n where
	{-# INLINE [0] suc' #-}
	suc' i s = do
	    step <- suc i s
	    case step of
	    	Yield (Box x) s' -> return (Yield x s')
	    	Last (Box x) -> return (Last x)
	    	Done -> return Done

{-# INLINE fromListN #-}
fromListN :: Monad m => Int -> [a] -> Stream m a
fromListN n xs = Stream suc xs n where
	suc _ (x:xs) = return (Yield x xs)
	suc _ [] = return Done