{-# LANGUAGE BangPatterns, MagicHash #-}

module Data.RangeMin.Cartesian (buildDepths, buildDepths1) where

import Control.Monad
import Data.RangeMin.Common.Types
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as UV
import qualified Data.Vector.Fusion.Stream.Monadic as SM
import Data.Vector.Fusion.Stream (Step (..), Stream)
import Data.RangeMin.Common.Vector
import Data.RangeMin.Cartesian.STInt
import Prelude hiding (read)

data IL a = IL {-# UNPACK #-} !Int a (IL a) | Nil

{-# INLINE mapAccumSM #-}
mapAccumSM :: Monad m => (b -> a -> m (c, b)) -> b -> SM.Stream m a -> SM.Stream m c
mapAccumSM f z0 (SM.Stream suc s0 n) = SM.Stream suc' (z0, s0) n where
	suc' (z, s) = do
		step <- suc s
		case step of
			Done	-> return Done
			Skip s'	-> return (Skip (z, s'))
			Yield x s' -> do
				(!y, z') <- f z x
				return (Yield y (z', s'))

{-# INLINE mapAccumS #-}
mapAccumS :: (b -> a -> (c, b)) -> b -> Stream a -> Stream c
mapAccumS f = mapAccumSM (\ b a -> return (f b a))

{-# INLINE buildDepths #-}
buildDepths :: G.Vector v a => LEq a -> v a -> UV.Vector Int
buildDepths (<=?) !xs = inlineCreate $ do
	!dest <- newWith n (-1)
	let depth !i = toSTInt $ do
		d0 <- read dest i
		case d0 of
		  -1 -> case parent i of
		    -1	-> do	write dest i 0
		    		return 0
		    p	-> do	!d' <- runSTInt $ depth p
		    		let !d = d' + 1
		    		write dest i d
		    		return $! d
		  _  -> return d0
	mapM_ (runSTInt . depth) [0..n-1]
	return dest
	where	!n = G.length xs
		lefts :: UV.Vector Int
		rights :: UV.Vector Int
		!lefts = unsafeVec n $ neighbors (<=?) (streamI xs)
		!rights = unsafeVec n $ neighbors (<=?) (streamIR xs)
		parent !i = case (lefts ! i, rights ! i) of
			(-1, -1)  -> -1
			(-1, r)   -> r
			(l, -1)   -> l
			(l, r)	  -> if (xs ! l) <=? (xs ! r) then r else l

buildDepths1 :: UV.Vector Int -> UV.Vector Int
buildDepths1 = buildDepths (<=)

{-# INLINE [1] neighbors #-}
neighbors :: LEq a -> Stream (Int, a) -> Stream IP
neighbors (<=?) xs = mapAccumS suc Nil xs
 where	suc stk (!i, xi) = let
	  run Nil = (IP i (-1), IL i xi Nil)
	  run stk0@(IL j xj stk)
		  | not (xi <=? xj)
			  = (IP i j, IL i xi stk0)
		  | otherwise
			  = run stk
	  in run stk