{-# LANGUAGE FlexibleContexts, BangPatterns #-}

-- | This is a specialized fusion framework designed for the particular needs
-- of this library.  In particular, it avoids duplicated indices, and has
-- special support for the 'snoc' operation.
-- 
-- It is also compatible with "Data.Vector.Fusion.Stream", being smart enough
-- to convert between the two styles as necessary.
-- 
-- Most methods in this module that output a vector return, in particular, a
-- 'VVector'.  If another vector implementation is required, the 'convert' function
-- will fuse away the intermediate 'VVector'.
module Data.RangeMin.Fusion where

import Control.Monad.ST (runST)
import Control.Monad.Primitive
import Data.RangeMin.Common.Types
import Data.RangeMin.Common.Vector
import Data.RangeMin.Common.Combinators
import qualified Data.Vector.Generic as G

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

import Prelude hiding (mapM_, map, replicate)

{-# INLINE convert #-}
convert :: (Vector v a, Vector v' a) => v a -> v' a
convert xs = unstream (stream xs)

{-# NOINLINE unstream #-}
unstream :: Vector v a => Stream a -> v a
unstream !str = create $ do
	!dest <- new (S.length str)
	fill dest str

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

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

{-# INLINE length #-}
length :: Vector v a => v a -> Length
length xs = S.length $! stream xs

{-# INLINE unzip #-}
unzip :: (Vector v (a, b), Vector va a, Vector vb b) =>
	v (a, b) -> (va a, vb b)
unzip xs = runST $ 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 :: (Vector v (a, b, c), Vector va a, Vector vb b, Vector vc c) =>
	v (a, b, c) -> (va a, vb b, vc c)
unzip3 xs = runST $ 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 :: Vector v a => v a -> Stream a
streamR !xs = fmap (G.unsafeIndex xs) (S.enumNR (G.length xs))

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

{-# INLINE enumN #-}
enumN :: Length -> VVector Int
enumN n = unstream (S.enumN n)

{-# INLINE enumNR #-}
enumNR :: Length -> VVector Int
enumNR n = unstream (S.enumNR n)

{-# INLINE generate #-}
generate :: Length -> (Index -> a) -> VVector a
generate n f = unstream (S.generate n f)

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

{-# INLINE map #-}
map :: (Vector v a) => (a -> a') -> v a -> VVector a'
map f = imap (const f)

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

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

{-# INLINE ipostscanl #-}
ipostscanl :: (Vector v a) => (b -> Index -> a -> b) -> b -> v a -> VVector b
ipostscanl f = imapAccumL (\ z i a -> let z' = f z i a in (z', z'))

{-# INLINE iterateN #-}
iterateN :: Length -> a -> (a -> a) -> VVector a
iterateN n z f = ipostscanl (\ z _ _ -> f z) z (replicate n ())

{-# INLINE mapAccumL #-}
mapAccumL :: (Vector v a) => (b -> a -> (c, b)) -> b -> v a -> VVector c
mapAccumL f = imapAccumL (\ z _ -> f z)

{-# INLINE unsafeUpdate #-}
unsafeUpdate :: (MVector v a, Vector v' (Index, a), PrimMonad m) => 
	v (PrimState m) a -> v' (Index, a) -> m ()
unsafeUpdate !dest = mapM_ (uncurry' (write dest))

{-# INLINE mapM_ #-}
mapM_ :: (Monad m, Vector v a) => (a -> m b) -> v a -> m ()
mapM_ f = imapM_ (const f)

{-# INLINE postscanl #-}
postscanl :: (Vector v a) => (b -> a -> b) -> b -> v a -> VVector b
postscanl f = ipostscanl (\ z _ -> f z)

{-# INLINE replicate #-}
replicate :: Int -> a -> VVector a
replicate n a = generate n (const a)

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

{-# INLINE snoc' #-}
snoc' :: Vector v a => v a -> a -> VVector a
xs `snoc'` x = x `seq` (xs `snoc` x)

{-# INLINE iunfoldN #-}
iunfoldN :: Length -> (Index -> b -> Maybe (a, b)) -> b -> VVector a
iunfoldN n f z = unstream (S.iunfoldN n f z)

{-# INLINE unfoldN #-}
unfoldN :: Int -> (b -> Maybe (a, b)) -> b -> VVector a
unfoldN n f = iunfoldN n (const f)

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

{-# INLINE foldl #-}
foldl :: Vector v a => (b -> a -> b) -> b -> v a -> b
foldl f = ifoldl (\ z _ -> f z)

{-# INLINE fromListN #-}
fromListN :: Length -> [a] -> VVector a
fromListN n xs = unstream (S.fromListN n xs)

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

{-# INLINE [0] munstream #-}
munstream :: (PrimMonad m, 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, Vector v a) => Length -> m a -> m (v a)
replicateM n m = munstream (SM.generateM n (const m))

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

{-# INLINE cap #-}
cap :: Vector v a => v a -> v a
cap xs = unstream (S.cap (stream xs))

{-# RULES
	"unstream/stream" forall xs . unstream (stream xs) = xs;
	"stream/unstream" forall str . stream (unstream str) = str;
-- 	"stream/Data.Vector.Generic.unstream"
-- 		forall str . stream (G.new (New.unstream str)) = S.fromVStream str;
-- 	"Data.Vector.Generic.stream/unstream"
-- 		forall str . G.stream (unstream str) = S.toVStream str;
	"unstream" [0] forall str . unstream str = create $ do
		!dest <- new $! S.length $! str
		S.imapM_ (write dest) (S.liftStream str)
		return dest;
	#-}