    {-# LANGUAGE BangPatterns, CPP, FlexibleContexts, Rank2Types #-}
    #if __GLASGOW_HASKELL__ >= 704
    {-# OPTIONS_GHC -fsimpl-tick-factor=200 #-}
    #endif

    -- |
    -- Module    : Statistics.Function
    -- Copyright : (c) 2009, 2010, 2011 Bryan O'Sullivan
    -- License   : BSD3
    --
    -- Maintainer  : bos@serpentine.com
    -- Stability   : experimental
    -- Portability : portable
    --
    -- Useful functions.

    module Statistics.Function
        (
        -- * Scanning
          minMax
        -- * Sorting
        , sort
        , gsort
        , sortBy
        , partialSort
        -- * Indexing
        , indexed
        , indices
        -- * Bit twiddling
        , nextHighestPowerOfTwo
        -- * Comparison
        , within
        -- * Arithmetic
        , square
        -- * Vectors
        , unsafeModify
        -- * Combinators
        , for
        , rfor
        ) where

    #include "MachDeps.h"

    import Control.Monad.ST (ST)
    import Data.Bits ((.|.), shiftR)
    import qualified Data.Vector.Algorithms.Intro as I
    import qualified Data.Vector.Generic as G
    import qualified Data.Vector.Unboxed as U
    import qualified Data.Vector.Unboxed.Mutable as M
    import Statistics.Function.Comparison (within)

    -- | Sort a vector.
    sort :: U.Vector Double -> U.Vector Double
    sort = G.modify I.sort
    {-# NOINLINE sort #-}

    -- | Sort a vector.
    gsort :: (Ord e, G.Vector v e) => v e -> v e
    gsort = G.modify I.sort
    {-# INLINE gsort #-}

    -- | Sort a vector using a custom ordering.
    sortBy :: (G.Vector v e) => I.Comparison e -> v e -> v e
    sortBy f = G.modify $ I.sortBy f
    {-# INLINE sortBy #-}

    -- | Partially sort a vector, such that the least /k/ elements will be
    -- at the front.
    partialSort :: (G.Vector v e, Ord e) =>
                   Int -- ^ The number /k/ of least elements.
                -> v e
                -> v e
    partialSort k = G.modify (`I.partialSort` k)
    {-# SPECIALIZE partialSort :: Int -> U.Vector Double -> U.Vector Double #-}

    -- | Return the indices of a vector.
    indices :: (G.Vector v a, G.Vector v Int) => v a -> v Int
    indices a = G.enumFromTo 0 (G.length a - 1)
    {-# INLINE indices #-}

    -- | Zip a vector with its indices.
    indexed :: (G.Vector v e, G.Vector v Int, G.Vector v (Int,e)) => v e -> v (Int,e)
    indexed a = G.zip (indices a) a
    {-# INLINE indexed #-}

    data MM = MM {-# UNPACK #-} !Double {-# UNPACK #-} !Double

    -- | Compute the minimum and maximum of a vector in one pass.
    minMax :: (G.Vector v Double) => v Double -> (Double, Double)
    minMax = fini . G.foldl' go (MM (1/0) (-1/0))
      where
        go (MM lo hi) k = MM (min lo k) (max hi k)
        fini (MM lo hi) = (lo, hi)
    {-# INLINE minMax #-}

    -- | Efficiently compute the next highest power of two for a
    -- non-negative integer.  If the given value is already a power of
    -- two, it is returned unchanged.  If negative, zero is returned.
    nextHighestPowerOfTwo :: Int -> Int
    nextHighestPowerOfTwo n
    #if WORD_SIZE_IN_BITS == 64
      = 1 + _i32
    #else
      = 1 + i16
    #endif
      where
        i0   = n - 1
        i1   = i0  .|. i0  `shiftR` 1
        i2   = i1  .|. i1  `shiftR` 2
        i4   = i2  .|. i2  `shiftR` 4
        i8   = i4  .|. i4  `shiftR` 8
        i16  = i8  .|. i8  `shiftR` 16
        _i32 = i16 .|. i16 `shiftR` 32
    -- It could be implemented as
    --
    -- > nextHighestPowerOfTwo n = 1 + foldl' go (n-1) [1, 2, 4, 8, 16, 32]
    --     where go m i = m .|. m `shiftR` i
    --
    -- But GHC do not inline foldl (probably because it's recursive) and
    -- as result function walks list of boxed ints. Hand rolled version
    -- uses unboxed arithmetic.

    -- | Multiply a number by itself.
    square :: Double -> Double
    square x = x * x

    -- | Simple for loop.  Counts from /start/ to /end/-1.
    for :: Monad m => Int -> Int -> (Int -> m ()) -> m ()
    for n0 !n f = loop n0
      where
        loop i | i == n    = return ()
               | otherwise = f i >> loop (i+1)
    {-# INLINE for #-}

    -- | Simple reverse-for loop.  Counts from /start/-1 to /end/ (which
    -- must be less than /start/).
    rfor :: Monad m => Int -> Int -> (Int -> m ()) -> m ()
    rfor n0 !n f = loop n0
      where
        loop i | i == n    = return ()
               | otherwise = let i' = i-1 in f i' >> loop i'
    {-# INLINE rfor #-}

    unsafeModify :: M.MVector s Double -> Int -> (Double -> Double) -> ST s ()
    unsafeModify v i f = do
      k <- M.unsafeRead v i
      M.unsafeWrite v i (f k)
    {-# INLINE unsafeModify #-}