{-# LANGUAGE FlexibleContexts, BangPatterns #-}

module Data.RangeMin.Common.Vector ((!), vec, unsafeVec, G.stream, streamRI, unsafeBackpermute', streamM, fillSlice,
	buildSliced) where

import Control.Monad.ST (ST)
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM
import qualified Data.Vector.Fusion.Stream as S
import qualified Data.Vector.Fusion.Stream.Monadic as SM
import Data.Vector.Fusion.Util
-- import Data.Vector.Fusion.Stream.Size

{-# INLINE streamM #-}
streamM :: (G.Vector v a, Monad m) => v a -> SM.Stream m a
streamM = S.liftStream . G.stream

(!) :: G.Vector v a => v a -> Int -> a
(!) = G.unsafeIndex

{-# INLINE fillSlice #-}
fillSlice :: G.Vector v a => Int -> Int -> G.Mutable v s a -> SM.Stream (ST s) a -> ST s (v a)
fillSlice !i !n !arr stream = do
	let slice = GM.unsafeSlice i n arr
	G.unsafeFreeze =<< GM.fill slice stream

{-# INLINE buildSliced #-}
buildSliced :: G.Vector v a => Int -> Int -> Maybe a -> v a ->
	(Int -> v a -> v a) -> v a
buildSliced !rows !cols !def initStream mkStream = G.create $ do
	let !size = rows * cols
	!arr <- maybe (GM.unsafeNew size) (GM.unsafeNewWith size) def
	!ini <- fillSlice 0 cols arr (streamM initStream)
	let filler !z !off !prev 
		| z < rows	= do
			!next <- fillSlice off cols arr (streamM (mkStream z prev))
			filler (z+1) (off+cols) next
		| otherwise	= return ()
	filler 1 cols ini
	return arr

{-# INLINE unsafeVec #-}
unsafeVec :: G.Vector v a => Int -> S.Stream (Int, a) -> v a
unsafeVec !n str = G.create $ do
	vec <- GM.unsafeNew n
	GM.unsafeUpdate vec str
	return vec

{-# INLINE vec #-}
vec :: G.Vector v a => Int -> S.Stream (Int, a) -> v a
vec !n str = G.create $ do
	vec <- GM.new n
	GM.update vec str
	return vec

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

streamRI :: G.Vector v a => v a -> S.Stream (Int, a)
{-# INLINE [1] streamRI #-}
streamRI !v = S.unfoldrN n get n
  where
    !n = G.length v
    {-# INLINE get #-}
    get 0 = Nothing
    get i = let !i' = i-1
            in
            case G.basicUnsafeIndexM v i' of Box x -> Just ((i', x), i')