{-# LANGUAGE MagicHash, BangPatterns, ParallelListComp, RankNTypes, ViewPatterns #-}
{-# OPTIONS_GHC -funbox-strict-fields #-}
{-|
This module is designed to implement high-speed versions of the function

@
rangeMinBy :: (Ix i, IArray a e) => (e -> e -> Ordering) -> a i e -> ((i, i) -> (i, e))
rangeMinBy cmp arr = \\r -> minimumBy (\\ (i1, e1) (i2, e2) -> cmp e1 e2) [(i, arr ! i) | i \<- range r]
@

for indices that are, in a sense, scalars -- that is, the array is one-dimensional.  This restriction is
encapsulated in the type restriction 'Enum'.

To be precise, the implementation this module supplies is O(n) and answers queries in constant time.

Based on analysis from "Test.Benchpress", making @q@ queries on a range-min of size @n@, 
constructed from an 'Int'-indexed list of 'Int's, takes approximately @0.0015n + 0.0023q@ ms on 
a Dell Latitude D630 with 2.50 GHz Intel Core Duo processors.

This module also includes implementations of a lowest-common-ancestor algorithm on trees with identical asymptotics.

Warning: the internal libraries of this module include several nonstandard list fusions.  For details, see
<http://hackage.haskell.org/packages/archive/rangemin/1.0.1/doc/html/src/Data-RangeMin-Internal-HandyList.html>.

Warning: This module does /not/ include bounds checking!
-}

-- An archaic version of this module took 1m 42-46 s for the same process.
module Data.RangeMin (
	-- * Convenient type aliases
	RangeMin, Comparator, 
	-- * Functions
	-- ** Range-min functions
	-- *** On IArrays
	rangeMinBy, rangeMin, rangeMinValue, rangeMinValueBy, rangeMinOpt,
	-- *** On lists
	rangeMinBy', rangeMin', rangeMinValue', rangeMinValueBy',
	-- *** On MArrays
	rangeMinByM, rangeMinM,
	-- ** LCA functions
	lowestCommonAncestor, lowestCommonAncestorBy, lowestCommonAncestorBy'
	 ) where

import Control.Monad (liftM, liftM2)
import Data.Maybe(fromJust, catMaybes, fromMaybe)
import Data.Tree (Tree(Node))
import Data.List(tails)
import Data.Monoid
import Data.Array.IArray(listArray, array, elems, bounds, Ix, Array, (!))
import GHC.Arr(unsafeIndex, unsafeRangeSize)
-- import Debug.Trace
import Data.Array.MArray(freeze)
import Data.Foldable (Foldable, toList)
import Data.Array.Base
import Data.Array.ST(STArray, runSTUArray)
import Control.Arrow((&&&))
import Data.Ix(Ix(range))
import Data.RangeMin.Internal.HandyArray
import Data.RangeMin.Internal.HandyOrdering (minimumBy, orderPair, minBy, comparing)
import Data.RangeMin.Internal.Combinators (onPair, on, (.:))
import Data.RangeMin.Internal.HandyNumeric
import Data.RangeMin.Internal.HandyList
import Control.Monad.ST
import Data.Array.Unboxed(UArray)
import GHC.Exts

-- | The type of a range-min function.  Takes as an argument a range of indices and returns
-- the position and value of the smallest element in that range.
type RangeMin i e = (i, i) -> (i, e)
	-- woo higher-order functions
type RangeMinVal i e = (i, i) -> e
	-- for internal use
type Comparator e = e -> e -> Ordering

{-# INLINE offset #-}
offset :: Enum i => i -> Int -> i
offset (fromEnum -> x) = toEnum . (x+)

{-# INLINE rangeSize #-}
rangeSize :: Enum i => (i, i) -> Int
rangeSize (fromEnum -> i, fromEnum -> j) = j-i+1

{-# INLINE index #-}
index :: Enum i => (i, i) -> i -> Int
index (fromEnum -> i, _) (fromEnum -> x) = x-i

{-# INLINE zeroEft' #-}
zeroEft' :: Int -> [Int]
zeroEft' i = zeroEft (i - 1)

--------------------------------------------------------------------

data IntAccum e = IA {-# UNPACK #-} !Int e -- hopefully unboxed by GHC
data BlockAccum e f = BA {-# UNPACK #-} !Int e f -- hopefully unboxed by GHC

-- folds/builds
{-# INLINE assocs' #-}
assocs' :: Enum i => (i, i) -> [e] -> [(i, e)]
assocs' (_, fromEnum -> iMax) l = {-# SCC "assocs" #-} build (\ c nil -> case foldr (assocs'' c) (IA iMax nil) l of IA _ ls -> ls) where
	assocs'' c x (IA i ls) = IA (i-1) ((toEnum i, x) `c` ls)

-- should fold/build, yay!
{-# INLINE blockify #-}
blockify :: Int -> Int -> [e] -> [(Int, [e])]
blockify n bS l =  build (\ c nil -> case foldr (blocker c) (BA n [] nil) l of BA _ _ bs -> bs) where
	!breakpoint = n - bS
	blocker c x (BA i b bs) = let j = i-1 in
		if j `rem` bS == 0 then BA j [] ((if j > breakpoint then n - j else bS, x:b) `c` bs) else BA j (x:b) bs

{-# INLINE boxer #-}
boxer :: (Int# -> Int#) -> Int -> Int
boxer f (I# i) = I# (f i)

{-# INLINE unboxer #-}
unboxer :: (Int -> Int) -> Int# -> Int#
unboxer f i = case f (I# i) of I# x -> x

{-# INLINE boxIn #-}
boxIn :: (Int# -> e) -> Int -> e
boxIn f (I# i) = f i

{-# INLINE unboxIn #-}
unboxIn :: (Int -> e) -> Int# -> e
unboxIn f i = f (I# i)

{--------------------------------------------------------------------
This module contains three implementations of rangeMinBy, and have the following asymptotics:
rangeMin3: O(n^2)
rangeMin2: O(n log n)
rangeMin1: O(n log* n)
(The asymptotics of rangeMin1, by the way, are achieved by bootstrapping to itself until
it reaches n small enough to confidently call up rangeMin2 or rangeMin3.)

While there is a strictly linear-time range-min algorithm, it requires converting
the input to a tree, traversing this tree in a special fashion, and calling a specialized
blocked range-min (see below for explanation) on the list of the depth reached at each step
of the traversal.  The time and memory overhead associated with this traversal
is such that it is considerably faster for any conceivable n to use this implementation.
-------------------------------------------------------------------}

{---------------------
O(n^2) range-min data structure with O(1) query time, with low constant factors on both construction 
and query.  Simply finds the minimum element in every subrange of the array.
----------------------}

{-# INLINE rangeMin3Look #-}
rangeMin3Look :: (Int# -> e) -> Comparator e -> Int -> [Int] -> RangeMinVal Int Int
rangeMin3Look (boxIn -> look) cmp !n ixs = {-# SCC "rangeMin3" #-}
	runST $ liftM (. uncurry encodePosition) $ mListUArray rmSize $ build (\ c n -> foldr (\ t -> scanl1FB min' t c) n (tails ixs))
	where	encodePosition !i !j = i * lastIx + j - half ((i - 1) * i)
		!rmSize = lastIx * n - half ((lastIx - 1) * lastIx) + 1 --encodePosition lastIx lastIx + 1
		!lastIx = n - 1
		min' = minBy (cmp `on` look)

{--------------------------------------------------------------------------------------------
O(n log n) construction, O(1) query time range-min algorithm, lazy in computation but does some allocation
immediately.  Relatively low constant factor.

The algorithm works by finding the minimum element of each range of width a power of two.  (Recursion 
allows each to be found in constant time from the previous level).  When a query is received of width n,
 let m be the largest power of two at most n; the two (possibly overlapping) ranges of width m that 
cover the query range are looked up and compared.
--------------------------------------------------------------------------------------------}

{-# INLINE rangeMin2Look #-}
rangeMin2Look :: (Int# -> e) -> Comparator e -> Int -> [Int] -> RangeMinVal Int Int
rangeMin2Look (boxIn -> look) cmp n ixs = {-# SCC "full_rM2" #-} \ r@(i, j) -> {-# SCC "rangeMin2_queries" #-}
	if i == j then initRow i else 
	let	!logWidth	= intLog (j - i)
		!left		= i
		!right		= j - pow2 logWidth + 1	in rM logWidth left right
	where	!logn		= intLog n + 1
		min'		= minBy (cmp `on` look)
		!(initRow, rMs) = runST $ 
			do	iR <- {-# SCC "initRow" #-} liftM unboxer $ mListUArray n ixs
				rm <- mListArray logn $ build (rower 1 iR)
				return (boxer iR, rm)
		nextRow !p !p' (boxer -> row) = 
			{-# SCC "rowBuilding" #-} runST $ liftM unboxer $ mFuncUArray (n - p' + 1) (\ i -> (min' `on` row) i (i+p))
		rower !p rowM c nil
			| p >= n	= nil
			| otherwise	= rowM `c` rower p' (nextRow p p' rowM) c nil
			where	!p' = p + p
		rM r		= min' `on` boxer (rMs r)
-- Note that no lists are being computed explicitly at all.   Woo!

{-----------------------------------------------------------------
Suppose we have a range-min algorithm that takes time f(n) for some superlinear f.  Then if 
the array is broken into blocks of size b, we can solve range-min queries as follows:

Break the query into a fractional piece of a block on each end with an integral number of 
blocks in the middle.  Given range-min functions for internal queries to each individual block,
 and for any integral range of blocks, we can then answer the query by making three separate 
subqueries, two internal block queries and one multi-block query.  Note that the range-min of 
an integral number of blocks is precisely a range-min of the array containing the minimum 
element of each block.  Furthermore, this algorithm will require only n/b f(b) + f(n/b) time.

For b=log n, and the f(n) = n log n construction, this is
n/(log n) log n log log n + n/log n log (n/log n) <= n log log n
In point of fact, replacing f(n) by this new n log log n gives an n log log log n construction, 
and so on, yielding a bootstrapped O(n log* n) construction.

   i  bi                bj  j
[--<--|-----|-----|-----|--->-|-]
----------------------------------------------------------------}

{-# INLINE blockedRangeMin #-}
blockedRangeMin :: Show e => Comparator e -> Int -> (Int -> RangeMinVal Int e) -> RangeMinVal Int e -> RangeMinVal Int e
blockedRangeMin (minBy -> min') !bS !blockRM !multiBlockRM (flip quotRem' bS -> (!bi,!xi), flip quotRem' bS -> (!bj, !xj)) =
	if bi == bj then {-# SCC "single_block_query" #-} blockRM bi (xi, xj) else
		let sidesMin = min' (blockRM bi (xi, lastIx)) (blockRM bj (0, xj)) in
				(if bi == pred bj then {-# SCC "two_block_query" #-} sidesMin else 
					{-# SCC "multi_block_query" #-} min' sidesMin (multiBlockRM (succ bi, pred bj)))
	where !lastIx = bS - 1

{-# INLINE rangeMin1Look #-}
rangeMin1Look :: (Int# -> e) -> Comparator e -> Int -> [Int] -> RangeMinVal Int Int
rangeMin1Look look cmp n ixs = 
	blockedRangeMin (cmp `on` boxIn look) blockSize blockRM multiBlockRM 
	where	!blockSize	= 1 + intLog n
		!nBlocks	= n `divCeil` blockSize
		blockRM		= {-# SCC "block_internal_rangemins" #-} runST $ mListArray nBlocks $
					map (uncurry (rangeMin3Look look cmp)) $ blockify n blockSize ixs
		bW		= blockWidth n blockSize
		multiBlockRM	= {-# SCC "multi_block_rangemins" #-} internalRangeMinBy look cmp nBlocks $
			map (\ i -> blockRM i (0, bW i - 1)) $ zeroEft' nBlocks

blockWidth :: Int -> Int -> Int -> Int
blockWidth !n !bS = let !nBlocks = n `divCeil` bS - 1 in \ !i -> if i == nBlocks then n `modCeil` bS else bS

---------------------------------------------------------------------------------

catalanArr :: UArray (Int, Int) Int
catalanArr = {-# SCC "catalanArr" #-} runSTUArray $ 
		do	!arr <- newArray_ r
			let writer !i !j = unsafeWrite arr (unsafeIndex r (i, j)) =<< cat arr i j
			mapM_ (\ !i -> mapM_ (writer i) $ unsafeEft i 16) $ zeroEft 16
			return arr
	where	r = ((0,0), (16,16))
		cat arr i j = case i of
			0	-> return 1
			1	-> return j
			_	-> liftM2 (+) (reader arr (i-1, j)) (reader arr (i, j-1))
		reader !arr !p = unsafeRead arr (unsafeIndex r p)

catalan :: (Int, Int) -> Int
catalan = unsafeLookup catalanArr

stackProcessor :: (e -> Bool) -> Int -> Int -> [e] -> (Int, Int, [e])
stackProcessor pr p = stackProc' 0 where
	stackProc' !cum !q stack = case stack of
		(x:xs)	| pr x	-> stackProc' (cum + catalan (p, q)) (q-1) xs
		_		-> (cum, q, stack)

data Ixr e = Ixr [e] {-# UNPACK #-} !Int {-# UNPACK #-} !Int {-# UNPACK #-} !Int

indexer :: Comparator e -> e -> Ixr e -> Ixr e
indexer cmp = {-# SCC "indexer" #-} index' where
	index' !x !(Ixr stack p q ans) = case stackProcessor (\ y -> cmp x y == LT) p q stack of
		(acc, q', stack') -> Ixr (x:stack') (p-1) q' (ans + acc)

{-# INLINE rangeMin0Look #-}
rangeMin0Look :: (Int# -> e) -> Comparator e -> Int -> [Int] -> RangeMinVal Int Int
rangeMin0Look look cmp n ixs =
	runST $	do	iLookup <- {-# SCC "iLookup" #-} liftM unboxer $ mListUArray n ixs
			indexRM <- newArr indexRange
			types <- {-# SCC "blockProcessing" #-} funcMUArray nBlocks (procBlock iLookup indexRM)
			indexRMIce <- mLook indexRM
			let blockRM !b = boxer iLookup . (blockSize * b +) . indexRMIce (types b)
			let blockMin !i = blockRM i (0, min (n - blockSize * i) blockSize - 1)--let !start = blockSize * i in minimumBy (cmp `on` boxIn look) $ map (boxer iLookup) (unsafeEft start (start + (min (n - start) blockSize - 1)))--
			return $ blockedRangeMin (cmp `on` boxIn look) blockSize blockRM $ internalRangeMinBy look cmp nBlocks (map blockMin  (zeroEft' nBlocks))
	where	!blockSize	= max 4 (ceilLog n `divCeil` 4)
		!nBlocks	= n `divCeil` blockSize
		indexRange	= catalan (blockSize, blockSize)
		look' iLookup	= boxIn look . boxer iLookup
		procBlock :: (Int# -> Int#) -> MArr (RangeMinVal Int Int) s -> Int -> ST s Int
		procBlock iLookup indexRM !b =
			let 	!start = blockSize * b
				!width = min blockSize (n - start) - 1
				lookIl = look' iLookup
				Ixr _ _ _ !ans = foldr (indexer cmp) (Ixr [] (blockSize-1) blockSize 0) $ map lookIl $ reverseEft start (start + width)
				in writeArr indexRM ans (rangeMin3Look I# (cmp `on` (lookIl . (+start))) (width+1) $ zeroEft width) >> return ans

---------------------------------------------------------------------------------
rangeMin3Size = 100
rangeMin2Size = 660
rangeMin1Size = 1200

rangeMinFunc :: Int -> a -> a -> a -> a -> a
rangeMinFunc !n f3 f2 f1 f0 = if n < rangeMin2Size then (if n < rangeMin3Size then f3 else f2) else (if n < rangeMin1Size then f1 else f0)

internalRangeMinBy :: (Int# -> e) -> Comparator e -> Int -> [Int] -> RangeMinVal Int Int
internalRangeMinBy look cmp !n l = {-# SCC "internal_RM" #-} let rM = lazyInternalRangeMinBy look cmp n l in rM (0, n-1) `seq` rM

lazyInternalRangeMinBy :: (Int# -> e) -> Comparator e -> Int -> [Int] -> RangeMinVal Int Int
lazyInternalRangeMinBy look cmp !n = rangeMinFunc n rangeMin3Look rangeMin2Look rangeMin1Look rangeMin0Look look cmp n

initRangeMin :: (Int# -> e) -> Comparator e -> Int -> RangeMinVal Int Int
initRangeMin look cmp n = rangeMin0Look look cmp n $ zeroEft' n

{-# SPECIALIZE rangeMinBy' :: Comparator e -> (Int, Int) -> [e] -> RangeMin Int e #-}
-- | Given a comparator, an index range and a list of elements in index order, returns a function that
-- takes two indices and returns the index and the value at the minimum value between those two indices inclusive.
rangeMinBy' :: (Enum i) => Comparator e	-- ^ A comparator function on elements.
					-> (i, i)	-- ^ Index range.
					-> [e]		-- ^ List of elements in index order.
					-> RangeMin i e	-- ^ Given a subrange, returns the index and value at the minimum in that range.
rangeMinBy' cmp bounds@(rangeSize -> n) elems = {-# SCC "initial_range_min" #-} 
	let (iMin, _) = bounds; look = runST $ liftM unboxIn $ mListArray n elems 
		in (offset iMin &&& boxIn look) . initRangeMin look cmp n . onPair (index bounds)
--rangeMinFunc n rangeMin3By rangeMin2By rangeMin1By cmp bounds

-- | A version of 'rangeMinBy'' on elements with their default ordering.
rangeMin' :: (Enum i, Ord e) => 
						(i, i) -- ^ Index range.
						-> [e] -- ^ List of elements in index order.
						-> RangeMin i e -- ^ Given a subrange, returns the index and value at the minimum in that range.
rangeMin' = rangeMinBy' compare

{-# SPECIALIZE rangeMinValueBy' :: Comparator e -> (Int, Int) -> [e] -> ((Int, Int) -> e) #-}
rangeMinValueBy' :: Enum i => Comparator e -> (i, i) -> [e] -> ((i, i) -> e)
rangeMinValueBy' cmp bounds@(rangeSize -> n) elems = 
	let look = runST $ liftM unboxIn $ mListArray n elems in boxIn look . initRangeMin look cmp n . onPair (index bounds)

rangeMinValue' :: (Enum i, Ord e) => (i, i) -> [e] -> ((i, i) -> e)
rangeMinValue' = rangeMinValueBy' compare

{-# SPECIALIZE rangeMinBy :: IArray a e => Comparator e -> a Int e -> RangeMin Int e #-}
-- | Given a comparator and an array, returns a function that takes a subrange of the array and
-- returns the index and value at the minimum in that subrange.
rangeMinBy :: (IArray a e, Enum i, Ix i) => Comparator e -- ^ A comparator function on elements.
					-> a i e -- ^ An array of elements.
					-> RangeMin i e -- ^ Given a subrange, returns the index and value at the minimum in that range.
rangeMinBy cmp arr@(numElements -> n) = let arrBounds@(iMin, _) = bounds arr; look = unboxIn $ unsafeAt arr in
	(offset iMin &&& boxIn look) . initRangeMin look cmp n . onPair (index arrBounds)

-- | A version of 'rangeMinBy' on elements with their default ordering.
rangeMin :: (IArray a e, Enum i, Ix i, Ord e) => a i e -> RangeMin i e
rangeMin = rangeMinBy compare

rangeMinOpt :: UArray Int Int -> RangeMinVal Int Int
rangeMinOpt arr@(numElements -> n) = unsafeAt arr . initRangeMin (unboxIn $ unsafeAt arr) compare n

{-# SPECIALIZE rangeMinValueBy :: Comparator e -> Array Int e -> ((Int, Int) -> e) #-}
{-# SPECIALIZE rangeMinValueBy :: Comparator Int -> UArray Int Int -> ((Int, Int) -> Int) #-}
rangeMinValueBy :: (IArray a e, Enum i, Ix i) => Comparator e -> a i e -> ((i, i) -> e)
rangeMinValueBy cmp !arr@(numElements -> n) = let arrBounds = bounds arr; look = unboxIn $ unsafeAt arr in
	boxIn look . initRangeMin look cmp n . onPair (index arrBounds)

rangeMinValue :: (IArray a e, Enum i, Ix i, Ord e) => a i e -> ((i, i) -> e)
rangeMinValue = rangeMinValueBy compare

rangeMinByM :: (Monad m, MArray a e m, Enum i, Ix i) => Comparator e -> a i e -> m (RangeMin i e)
rangeMinByM cmp = liftM (rangeMinBy cmp . asPureArray) . freeze

rangeMinM :: (Monad m, MArray a e m, Enum i, Ix i, Ord e) => a i e -> m (RangeMin i e)
rangeMinM = rangeMinByM compare

-----------------------------------------------------------

-- Traverses the elements of the tree in an order such that the lowest common ancestor of two nodes is always the
-- minimum element between them.  Returns a list of depth and node-label pairs as well as the list's length.
lcaTraversal :: Tree e -> (Int, [(Int, e)])
lcaTraversal t = case lcaTraversal' 0 t (IA 0 []) of IA n l -> (n, l); where
	lcaTraversal' !d (Node a ts) trav = let cat = lcaTraversal' (d+1); consHere (IA m tr) = IA (m+1) ((d, a):tr) in case ts of
		[]	-> consHere trav
		[t]	-> consHere (t `cat` trav)
		(t:ts) -> t `cat` foldr (consHere .: cat) trav ts

-- | Given an indexing function on node labels, returns a function that takes two node 
-- labels and returns the label of their lowest common ancestor.
lowestCommonAncestorBy :: Ix i => 	   (a -> Maybe i) -- ^ An index on tree labels so they can be indexed on an array. 
							--  Need not account for every node label, hence the Maybe.
					-> (i, i)	-- ^ Bounds on indices of tree labels.
					-> Tree a	-- ^ The tree to provide lowest-common-ancestor service for.
					-> ((a, a) -> a)	-- ^ Given two node labels, returns the label of their lowest common ancestor.
lowestCommonAncestorBy ix r tree = lowestCommonAncestorBy' ix r tree . onPair (fromJust . ix)

-- | Given an indexing function on node labels, returns a function that takes two indices and returns the label of the associated nodes' lowest common ancestor.
lowestCommonAncestorBy' :: Ix i => (a -> Maybe i) -- ^ An index on tree labels, so they can be indexed in an array.
							-- Need not account for every node label, hence the Maybe.
				-> (i, i)	-- ^ Bounds on indices of tree labels.
				-> Tree a	-- ^ The tree to provide lowest-common-ancestor service for.
				-> ((i, i) -> a)	-- ^ Given two indices, return the label of the associated nodes' lowest common ancestor.
							-- The behavior of this method is not specified if these indices do not correspond
							-- to any trees.
lowestCommonAncestorBy' ix r (lcaTraversal -> (n, trav)) = snd . rM . orderPair . onPair indexes
	where	pairList	= [(x, i) | (i, ix . snd -> Just x) <- zipNaturals n trav]
		rM		= rangeMinValueBy' (comparing fst) (0, n-1) trav
		indexes		= unsafeLookup $ asPureArray $ array r pairList

-- | Given a tree of indexed nodes, returns a function that takes two nodes and returns their lowest common ancestor.
lowestCommonAncestor :: Ix i => (i, i) -- ^ Bounds on node labels.
				-> Tree i -- ^ The tree to provide lowest-common-ancestor service for.
				-> ((i, i) -> i) -- ^ A function that takes two node labels and returns the label of their lowest common ancestor.
lowestCommonAncestor = lowestCommonAncestorBy Just

{-# RULES
	"fromEnum/Int" forall (x :: Int) . fromEnum x = x;
	"toEnum/Int" forall (x :: Int) . toEnum x = x;
	#-}