{-# LANGUAGE FlexibleContexts, BangPatterns #-}

module Data.RangeMin.Fusion where

import Control.Monad
import Control.Monad.Primitive
import Data.RangeMin.Common.ST
import Data.RangeMin.Common.Vector
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM

import Data.RangeMin.Fusion.Stream
import qualified Data.RangeMin.Fusion.Stream as S
import qualified Data.RangeMin.Fusion.Stream.Monadic as SM

{-# INLINE [0] unstream #-}
unstream :: G.Vector v a => Stream a -> v a
unstream str = inlineCreate $ do
	!dest <- new $! S.length str
	fill dest str

{-# INLINE [1] stream #-}
stream :: G.Vector v a => v a -> Stream a
stream !src = S.generate (G.length src) (G.unsafeIndex src)

{-# INLINE [0] istream #-}
istream :: G.Vector v a => v a -> Stream (Int, a)
istream !src = S.generate (G.length src) (\ i -> (i, G.unsafeIndex src i))

{-# INLINE unzip #-}
unzip :: (G.Vector v (a, b), G.Vector va a, G.Vector vb b) =>
	v (a, b) -> (va a, vb b)
unzip xs = inlineRunST $ do
	!destA <- new n
	!destB <- new n
	S.imapM_ (\ i (a, b) -> do
		write destA i a
		write destB i b) str
	liftM2 (,) (unsafeFreeze destA) (unsafeFreeze destB)
	where	str = stream xs
		n = S.length str

{-# INLINE unzip3 #-}
unzip3 :: (G.Vector v (a, b, c), G.Vector va a, G.Vector vb b, G.Vector vc c) =>
	v (a, b, c) -> (va a, vb b, vc c)
unzip3 xs = inlineRunST $ do
	!destA <- new n
	!destB <- new n
	!destC <- new n
	S.imapM_ (\ i (a, b, c) -> do
		write destA i a
		write destB i b
		write destC i c) str
	liftM3 (,,) (unsafeFreeze destA) (unsafeFreeze destB) (unsafeFreeze destC)
	where	str = stream xs
		n = S.length str

{-# INLINE streamR #-}
streamR :: G.Vector v a => v a -> Stream a
streamR !xs = S.map (G.unsafeIndex xs) (S.enumNR (G.length xs))

{-# INLINE streamIR #-}
streamIR :: G.Vector v a => v a -> Stream (Int, a)
streamIR !xs = S.map (\ i -> (i, G.unsafeIndex xs i)) (S.enumNR (G.length xs))

{-# INLINE enumN #-}
enumN :: G.Vector v Int => Int -> v Int
enumN n = unstream (S.enumN n)

{-# INLINE enumNR #-}
enumNR :: G.Vector v Int => Int -> v Int
enumNR n = unstream (S.enumNR n)

{-# INLINE generate #-}
generate :: G.Vector v a => Int -> (Int -> a) -> v a
generate n f = unstream (S.generate n f)

{-# INLINE imap #-}
imap :: (G.Vector v a, G.Vector v' a') => (Int -> a -> a') -> v a -> v' a'
imap f xs = unstream (S.imap f (stream xs))

{-# INLINE map #-}
map :: (G.Vector v a, G.Vector v' a') => (a -> a') -> v a -> v' a'
map f xs = unstream (S.map f (stream xs))

{-# INLINE imapAccumL #-}
imapAccumL :: (G.Vector v a, G.Vector v' c) => (b -> Int -> a -> (c, b)) -> b -> v a -> v' c
imapAccumL f z xs = unstream (S.imapAccumL f z (stream xs))

{-# INLINE imapM_ #-}
imapM_ :: (Monad m, G.Vector v a) => (Int -> a -> m b) -> v a -> m ()
imapM_ f xs = S.imapM_ f (stream xs)

{-# INLINE ipostscanl #-}
ipostscanl :: (G.Vector v a, G.Vector v' b) => (b -> Int -> a -> b) -> b -> v a -> v' b
ipostscanl f z xs = unstream (S.ipostscanl f z (stream xs))

{-# INLINE iterateN #-}
iterateN :: G.Vector v a => Int -> a -> (a -> a) -> v a
iterateN n z f = unstream (S.iterateN n z f)

{-# INLINE izipWith #-}
izipWith :: (G.Vector v1 a, G.Vector v2 b, G.Vector v3 c) => 
	(Int -> a -> b -> c) -> v1 a -> v2 b -> v3 c
izipWith f xs ys = unstream (S.izipWith f (stream xs) (stream ys))

{-# INLINE mapAccumL #-}
mapAccumL :: (G.Vector v a, G.Vector v' c) => (b -> a -> (c, b)) -> b -> v a -> v' c
mapAccumL f z xs = unstream (S.mapAccumL f z (stream xs))

{-# INLINE unsafeUpdate #-}
unsafeUpdate :: (GM.MVector v a, G.Vector v' (Int, a), PrimMonad m) => 
	v (PrimState m) a -> v' (Int, a) -> m ()
unsafeUpdate !dest src = S.mapM_ (\ (i, a) -> write dest i a) (stream src)

{-# INLINE mapM_ #-}
mapM_ :: (Monad m, G.Vector v a) => (a -> m b) -> v a -> m ()
mapM_ f xs = S.mapM_ f (stream xs)

{-# INLINE postscanl #-}
postscanl :: (G.Vector v a, G.Vector v' b) => (b -> a -> b) -> b -> v a -> v' b
postscanl f z xs = unstream (S.postscanl f z (stream xs))

{-# INLINE replicate #-}
replicate :: G.Vector v a => Int -> a -> v a
replicate n a = unstream (S.replicate n a)

{-# INLINE snoc #-}
snoc :: G.Vector v a => v a -> a -> v a
xs `snoc` x = unstream (stream xs `S.snoc` x)

{-# INLINE zipWith #-}
zipWith :: (G.Vector v1 a, G.Vector v2 b, G.Vector v3 c) =>
	(a -> b -> c) -> v1 a -> v2 b -> v3 c
zipWith f xs ys = unstream (S.zipWith f (stream xs) (stream ys))

{-# INLINE iunfoldN #-}
iunfoldN :: G.Vector v a => Int -> (Int -> b -> Maybe (a, b)) -> b -> v a
iunfoldN n f z = unstream (S.iunfoldN n f z)

{-# INLINE unfoldN #-}
unfoldN :: G.Vector v a => Int -> (b -> Maybe (a, b)) -> b -> v a
unfoldN n f z = unstream (S.unfoldN n f z)

{-# INLINE ifoldl #-}
ifoldl :: G.Vector v a => (b -> Int -> a -> b) -> b -> v a -> b
ifoldl f z xs = S.ifoldl f z (stream xs)

{-# INLINE foldl #-}
foldl :: G.Vector v a => (b -> a -> b) -> b -> v a -> b
foldl f z xs = S.foldl f z (stream xs)

{-# INLINE fromListN #-}
fromListN :: G.Vector v a => Int -> [a] -> v a
fromListN n xs = unstream (S.fromListN n xs)

{-# INLINE unsafeBackpermute #-}
unsafeBackpermute :: (G.Vector v a, G.Vector v' Int) => v a -> v' Int -> v a
unsafeBackpermute !xs ixs = 
	unstream (S.unbox (S.map (G.unsafeIndexM xs) (stream ixs)))

{-# INLINE [0] munstream #-}
munstream :: (PrimMonad m, G.Vector v a) => S.MStream m a -> m (v a)
munstream str = do
	let !n = SM.length str
	!dest <- new n
	_ <- fillM dest str
	unsafeFreeze (sliceM 0 n dest)

{-# INLINE [0] replicateM #-}
replicateM :: (PrimMonad m, G.Vector v a) => Int -> m a -> m (v a)
replicateM n m = munstream (SM.replicateM n m)

{-# INLINE [0] fillM #-}
fillM :: (PrimMonad m, GM.MVector v a) => v (PrimState m) a -> S.MStream m a -> m (v (PrimState m) a)
fillM !dest !str = do
	let !n = SM.length str
	SM.imapM_ (write dest) str
	return (sliceM 0 n dest)

{-# INLINE [0] fill #-}
fill :: (PrimMonad m, GM.MVector v a) => v (PrimState m) a -> Stream a -> m (v (PrimState m) a)
fill dest str = fillM dest (S.liftStream str)

{-# RULES
	"unstream/stream" forall xs . unstream (stream xs) = xs;
	"stream/unstream" forall str . stream (unstream str) = str;
	"unstream" [0] forall str . unstream str = inlineCreate $ do
		!dest <- new (S.length str)
		S.imapM_ (write dest) (S.liftStream str)
		return dest;
	#-}