module Select (select) where

import Control.Monad(Monad(..), when)
import Data.Array.MArray(readArray, writeArray, MArray, Ix)

select :: (Integral i, Ix i, Ord e, MArray a e m) =>
	  a i e -> i -> i -> i -> m()
select x l' r' k
    = whileM (uncurry (>)) select1 (r',l') >> return ()
    where
     select1 rl
	 = select2 rl >> select3 rl
     select2 (r, l)
	 = when ((r - l) > 600) $ do
	      let n  = fromIntegral $ r - l + 1
		  i  = fromIntegral $ k - l + 1
		  z  = log n
		  s  = 0.5 * exp (2 * z / 3)
		  sd = 0.5 * (sqrt (z * s * (((n - s) / n)))) * (signum (i - n / 2))
		  ll = max l (floor ((fromIntegral k) - i * s / n + sd))
		  rr = min r (floor ((fromIntegral k) + (n - i) * s / n + sd))
	      select x ll rr k
     select3 (r, l)
	 = do t <- readArray x k
	      let i = l
		  j = r
	      exchange x l k
	      (readArray x r) >>= \xr -> when (xr > t) (exchange x r l)
	      (i0, j0) <- whileM (uncurry (<)) (uncurry (select4 t)) (i,j)
	      let i = i0
		  j = j0
	      xl <- readArray x l
	      j0 <- if xl == t
		       then exchange x l j >> return j
		       else exchange x (j+1) r >> return (j+1)
	      let l0 | j0 <= k   = j0 + 1
		     | otherwise = l
		  r0 | k <= j0   = j0 - 1
		     | otherwise = r
	      return (r0,l0)	      
     select4 t i j
	 = do exchange x i j
	      let i0 = i + 1
		  j0 = j - 1
	      i1 <- select40 i0
	      j1 <- select41 j0
	      return (i1, j1)
	 where
	  select40 i0 
	      = do xi <- readArray x i0
		   if xi < t 
		      then select40 (i0 + 1)
		      else return i0
	  select41 j0
	      = do xj <- readArray x j0
		   if xj > t 
		      then select41 (j0 - 1)
		      else return j0

exchange seq0 i j 
    = do e0 <- readArray seq0 i
	 e1 <- readArray seq0 j
	 writeArray seq0 j e0
	 writeArray seq0 i e1

whileM :: (Monad m) => (a -> Bool) -> (a -> m a) -> a -> m a
whileM p f a
    | p a       = f a >>= whileM p f
    | otherwise = return a