    {-# LANGUAGE PatternGuards #-}

    module Numeric.Histogram ( Range
                             , binBounds
                             , histValues
                             , histWeightedValues
                             , histWithBins
                             ) where

    import qualified Data.Vector as V
    import qualified Data.Vector.Mutable as MV
    import Control.Monad.ST

    type Range a = (a,a)

    -- | 'binBounds a b n' generates bounds for 'n' bins spaced linearly between
    -- 'a' and 'b'
    --
    -- Examples:
    --
    -- >>> binBounds 0 3 4
    -- [(0.0,0.75),(0.75,1.5),(1.5,2.25),(2.25,3.0)]
    binBounds :: RealFrac a => a -> a -> Int -> [Range a]
    binBounds :: a -> a -> Int -> [Range a]
    binBounds a :: a
    a b :: a
    b n :: Int
    n = (Int -> Range a) -> [Int] -> [Range a]
    forall a b. (a -> b) -> [a] -> [b]
    map (\i :: Int
    i->(Int -> a
    forall a. Real a => a -> a
    lbound Int
    i, Int -> a
    forall a. Real a => a -> a
    lbound (Int
    iInt -> Int -> Int
    forall a. Num a => a -> a -> a
    +1))) [0..Int
    nInt -> Int -> Int
    forall a. Num a => a -> a -> a
    -1]
            where lbound :: a -> a
    lbound i :: a
    i = a
    a a -> a -> a
    forall a. Num a => a -> a -> a
    + (a
    ba -> a -> a
    forall a. Num a => a -> a -> a
    -a
    a) a -> a -> a
    forall a. Num a => a -> a -> a
    * a -> a
    forall a b. (Real a, Fractional b) => a -> b
    realToFrac a
    i a -> a -> a
    forall a. Fractional a => a -> a -> a
    / Int -> a
    forall a b. (Real a, Fractional b) => a -> b
    realToFrac Int
    n

    -- | 'histValues a b n vs' returns the bins for the histogram of
    -- 'vs' on the range from 'a' to 'b' with 'n' bins
    histValues :: RealFrac a => a -> a -> Int -> [a] -> V.Vector (Range a, Int)
    histValues :: a -> a -> Int -> [a] -> Vector (Range a, Int)
    histValues a :: a
    a b :: a
    b n :: Int
    n = Vector (Range a) -> [(Int, a)] -> Vector (Range a, Int)
    forall w a.
    (Num w, RealFrac a) =>
    Vector (Range a) -> [(w, a)] -> Vector (Range a, w)
    histWithBins ([Range a] -> Vector (Range a)
    forall a. [a] -> Vector a
    V.fromList ([Range a] -> Vector (Range a)) -> [Range a] -> Vector (Range a)
    forall a b. (a -> b) -> a -> b
    $ a -> a -> Int -> [Range a]
    forall a. RealFrac a => a -> a -> Int -> [Range a]
    binBounds a
    a a
    b Int
    n) ([(Int, a)] -> Vector (Range a, Int))
    -> ([a] -> [(Int, a)]) -> [a] -> Vector (Range a, Int)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [Int] -> [a] -> [(Int, a)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (Int -> [Int]
    forall a. a -> [a]
    repeat 1)

    -- | 'histValues a b n vs' returns the bins for the weighted histogram of
    -- 'vs' on the range from 'a' to 'b' with 'n' bins
    histWeightedValues :: RealFrac a => a -> a -> Int -> [(Double,a)] -> V.Vector (Range a, Double)
    histWeightedValues :: a -> a -> Int -> [(Double, a)] -> Vector (Range a, Double)
    histWeightedValues a :: a
    a b :: a
    b n :: Int
    n = Vector (Range a) -> [(Double, a)] -> Vector (Range a, Double)
    forall w a.
    (Num w, RealFrac a) =>
    Vector (Range a) -> [(w, a)] -> Vector (Range a, w)
    histWithBins ([Range a] -> Vector (Range a)
    forall a. [a] -> Vector a
    V.fromList ([Range a] -> Vector (Range a)) -> [Range a] -> Vector (Range a)
    forall a b. (a -> b) -> a -> b
    $ a -> a -> Int -> [Range a]
    forall a. RealFrac a => a -> a -> Int -> [Range a]
    binBounds a
    a a
    b Int
    n)

    -- | 'histWithBins bins xs' is the histogram of weighted values 'xs' with 'bins'
    --
    -- Examples:
    --
    -- >>> :{
    -- histWithBins
    --     (V.fromList [(0.0, 0.75), (0.75, 1.5), (1.5, 2.25), (2.25, 3.0)])
    --     [(1, 0), (1, 0), (1, 1), (1, 2), (1, 2), (1, 2), (1, 3)]
    -- :}
    -- [((0.0,0.75),2),((0.75,1.5),1),((1.5,2.25),3),((2.25,3.0),1)]
    histWithBins :: (Num w, RealFrac a) => V.Vector (Range a) -> [(w, a)] -> V.Vector (Range a, w)
    histWithBins :: Vector (Range a) -> [(w, a)] -> Vector (Range a, w)
    histWithBins bins :: Vector (Range a)
    bins xs :: [(w, a)]
    xs =
        let n :: Int
    n = Vector (Range a) -> Int
    forall a. Vector a -> Int
    V.length Vector (Range a)
    bins
            testBin :: RealFrac a => a -> (Int, Range a) -> Bool
            testBin :: a -> (Int, Range a) -> Bool
    testBin x :: a
    x (i :: Int
    i, (a :: a
    a,b :: a
    b)) =
                if Int
    i Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    n Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - 1
                    then a
    x a -> a -> Bool
    forall a. Ord a => a -> a -> Bool
    >= a
    a Bool -> Bool -> Bool
    && a
    x a -> a -> Bool
    forall a. Ord a => a -> a -> Bool
    <= a
    b
                    else a
    x a -> a -> Bool
    forall a. Ord a => a -> a -> Bool
    >= a
    a Bool -> Bool -> Bool
    && a
    x a -> a -> Bool
    forall a. Ord a => a -> a -> Bool
    < a
    b

            f :: (RealFrac a, Num w)
              => V.Vector (Range a) -> MV.STVector s w -> (w, a)
              -> ST s ()
            f :: Vector (Range a) -> STVector s w -> (w, a) -> ST s ()
    f bins1 :: Vector (Range a)
    bins1 bs :: STVector s w
    bs (w :: w
    w,x :: a
    x) =
                case ((Int, Range a) -> Bool)
    -> Vector (Int, Range a) -> Vector (Int, Range a)
    forall a. (a -> Bool) -> Vector a -> Vector a
    V.dropWhile (Bool -> Bool
    not (Bool -> Bool)
    -> ((Int, Range a) -> Bool) -> (Int, Range a) -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . a -> (Int, Range a) -> Bool
    forall a. RealFrac a => a -> (Int, Range a) -> Bool
    testBin a
    x) (Vector (Int, Range a) -> Vector (Int, Range a))
    -> Vector (Int, Range a) -> Vector (Int, Range a)
    forall a b. (a -> b) -> a -> b
    $ Vector (Range a) -> Vector (Int, Range a)
    forall a. Vector a -> Vector (Int, a)
    V.indexed Vector (Range a)
    bins1 of
                    v :: Vector (Int, Range a)
    v | Vector (Int, Range a) -> Bool
    forall a. Vector a -> Bool
    V.null Vector (Int, Range a)
    v  -> () -> ST s ()
    forall (m :: * -> *) a. Monad m => a -> m a
    return ()
                    v :: Vector (Int, Range a)
    v | (idx :: Int
    idx,_) <- Vector (Int, Range a) -> (Int, Range a)
    forall a. Vector a -> a
    V.head Vector (Int, Range a)
    v  -> do
                        w
    m <- MVector (PrimState (ST s)) w -> Int -> ST s w
    forall (m :: * -> *) a.
    PrimMonad m =>
    MVector (PrimState m) a -> Int -> m a
    MV.read STVector s w
    MVector (PrimState (ST s)) w
    bs Int
    idx
                        MVector (PrimState (ST s)) w -> Int -> w -> ST s ()
    forall (m :: * -> *) a.
    PrimMonad m =>
    MVector (PrimState m) a -> Int -> a -> m ()
    MV.write STVector s w
    MVector (PrimState (ST s)) w
    bs Int
    idx (w -> ST s ()) -> w -> ST s ()
    forall a b. (a -> b) -> a -> b
    $! w
    mw -> w -> w
    forall a. Num a => a -> a -> a
    +w
    w

            counts :: Vector w
    counts = (forall s. ST s (Vector w)) -> Vector w
    forall a. (forall s. ST s a) -> a
    runST ((forall s. ST s (Vector w)) -> Vector w)
    -> (forall s. ST s (Vector w)) -> Vector w
    forall a b. (a -> b) -> a -> b
    $ do STVector s w
    b <- Int -> w -> ST s (MVector (PrimState (ST s)) w)
    forall (m :: * -> *) a.
    PrimMonad m =>
    Int -> a -> m (MVector (PrimState m) a)
    MV.replicate Int
    n 0
                                ((w, a) -> ST s ()) -> [(w, a)] -> ST s ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (Vector (Range a) -> STVector s w -> (w, a) -> ST s ()
    forall a w s.
    (RealFrac a, Num w) =>
    Vector (Range a) -> STVector s w -> (w, a) -> ST s ()
    f Vector (Range a)
    bins STVector s w
    b) [(w, a)]
    xs
                                MVector (PrimState (ST s)) w -> ST s (Vector w)
    forall (m :: * -> *) a.
    PrimMonad m =>
    MVector (PrimState m) a -> m (Vector a)
    V.freeze STVector s w
    MVector (PrimState (ST s)) w
    b
        in Vector (Range a) -> Vector w -> Vector (Range a, w)
    forall a b. Vector a -> Vector b -> Vector (a, b)
    V.zip Vector (Range a)
    bins Vector w
    counts