{-# LANGUAGE Rank2Types, FlexibleContexts, BangPatterns #-}

module Data.RangeMin.Common.Vector (vec, unsafeVec, unsafeBackpermute', enumFromN, foldlRange, replicateM, inlineCreate,
	inlineUnstream, inlineUnstreamR, inlineBuild, hintSize, module Data.RangeMin.Common.Vector.Utils,
	streamI, streamIR) where

import Control.Monad hiding (replicateM)
import Data.RangeMin.Common.ST
import Data.RangeMin.Common.Types
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM
import qualified Data.Vector.Generic.New as New
import qualified Data.Vector.Fusion.Stream as S
import Data.RangeMin.Common.Vector.Utils
import Data.Vector.Fusion.Stream.Size
import Control.Monad.Primitive
import Prelude hiding (drop, read)

{-# INLINE hintSize #-}
hintSize :: Vector v a => Int -> v a -> v a
hintSize n xs = G.unstream (G.stream xs `S.sized` Exact n)

{-# INLINE [1] inlineCreate #-}
inlineCreate :: Vector v a => (forall s . ST s (Mutable v s a)) -> v a
inlineCreate m = inlineNew (New.New m)

{-# INLINE [1] replicateM #-}
replicateM :: (Vector v a, PrimMonad m) => Int -> m a -> m (v a)
replicateM !n x = do
	!arr <- new n
	forM_ [0..n-1] $ \ i -> write arr i =<< x
	unsafeFreeze arr

{-# INLINE foldlRange #-}
foldlRange :: (Int -> Int -> Int) -> Int -> Int -> Int
foldlRange f !i !n = S.foldl1' f (enumFromN i n)

{-# INLINE unsafeVec #-}
unsafeVec :: Vector v Int => Int -> S.Stream IP -> v Int
unsafeVec !n str = inlineCreate $ do
	!vec <- new n
	S.mapM_ (\ (IP i x) -> write vec i x) str
	return vec

{-# INLINE enumFromN #-}
enumFromN :: Int -> Int -> S.Stream Int
enumFromN !i !n = S.unfoldr suc i `S.sized` Exact n where
	!end = i + n
	suc i	| i >= end	= Nothing
		| otherwise	= Just (i, i+1)

{-# INLINE inlineUnstream #-}
inlineUnstream :: Vector v a => S.Stream a -> v a
inlineUnstream s = inlineNew (New.unstream s) 

{-# INLINE inlineUnstreamR #-}
inlineUnstreamR :: Vector v a => S.Stream a -> v a
inlineUnstreamR s = inlineNew (New.unstreamR s)

{-# INLINE inlineNew #-}
inlineNew :: Vector v a => New.New v a -> v a
inlineNew !m = inlineRunST (unsafeFreeze =<< New.run m)

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

{-# INLINE streamIR #-}
streamIR :: Vector v a => v a -> S.Stream (Int, a)
streamIR !xs = S.unfoldr suc n `S.sized` Exact n where
	!n = G.length xs
	suc i = if i == 0 then Nothing else let
	  !i' = i - 1
	  xi' = xs ! i'
	  in Just ((i', xi'), i')

{-# INLINE inlineBuild #-}
inlineBuild :: Vector v a => v a -> v a
inlineBuild xs = inlineUnstream (G.stream xs)

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

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