{-# LANGUAGE ViewPatterns #-}

{-# OPTIONS_GHC -fno-warn-unused-imports #-}

-- |

module ArrayFire.Statistics where

import ArrayFire.Array

import ArrayFire.FFI

import ArrayFire.Internal.Statistics

import ArrayFire.Internal.Types

-- | Calculates 'mean' of 'Array' along user-specified dimension.

mean

:: AFType a

=> Array a

-- ^ Input 'Array'

-> Int

-- ^ The dimension along which the mean is extracted

-> Array a

-- ^ Will contain the mean of the input 'Array' along dimension dim

mean a n =

a `op1` (\x y ->

af_mean x y (fromIntegral n))

-- | Calculates 'meanWeighted' of 'Array' along user-specified dimension.

meanWeighted

:: AFType a

=> Array a

-- ^ Input 'Array'

-> Array a

-- ^ Weights 'Array'

-> Int

-- ^ The dimension along which the mean is extracted

-> Array a

-- ^ Will contain the mean of the input 'Array' along dimension dim

meanWeighted x y (fromIntegral -> n) =

op2 x y $ \a b c ->

af_mean_weighted a b c n

-- | Calculates /variance/ of 'Array' along user-specified dimension.

var

:: AFType a

=> Array a

-- ^ Input 'Array'

-> Bool

-- ^ boolean denoting Population variance (false) or Sample Variance (true)

-> Int

-- ^ The dimension along which the variance is extracted

-> Array a

-- ^ will contain the variance of the input array along dimension dim

var arr (fromIntegral . fromEnum -> b) d =

arr `op1` (\p x ->

af_var p x b (fromIntegral d))

-- | Calculates 'varWeighted' of 'Array' along user-specified dimension.

varWeighted

:: AFType a

=> Array a

-- ^ Input 'Array'

-> Array a

-- ^ Weights 'Array' used to scale input in before getting variance

-> Int

-- ^ The dimension along which the variance is extracted

-> Array a

-- ^ Contains the variance of the input array along dimension dim

varWeighted x y (fromIntegral -> n) =

op2 x y $ \a b c ->

af_var_weighted a b c n

-- | Calculates 'stdev' of 'Array' along user-specified dimension.

stdev

:: AFType a

=> Array a

-- ^ Input 'Array'

-> Int

-- ^ The dimension along which the standard deviation is extracted

-> Array a

-- ^ Contains the standard deviation of the input array along dimension dim

stdev a n =

a `op1` (\x y ->

af_stdev x y (fromIntegral n))

-- | Calculates /covariance/ of two 'Array's with a bias specifier.

cov

:: AFType a

=> Array a

-- ^ First input 'Array'

-> Array a

-- ^ Second input 'Array'

-> Bool

-- ^ A boolean specifying if biased estimate should be taken (default: 'False')

-> Array a

-- ^ Contains will the covariance of the input 'Array's

cov x y (fromIntegral . fromEnum -> n) =

op2 x y $ \a b c ->

af_cov a b c n

-- | Calculates 'median' of 'Array' along user-specified dimension.

median

:: AFType a

=> Array a

-- ^ Input 'Array'

-> Int

-- ^ Dimension along which to calculate 'median'

-> Array a

-- ^ Array containing 'median'

median a n =

a `op1` (\x y ->

af_median x y (fromIntegral n))

-- | Calculates 'mean' of all elements in an 'Array'

meanAll

:: AFType a

=> Array a

-- ^ Input 'Array'

-> (Double, Double)

-- ^ Mean result (real and imaginary part)

meanAll = (`infoFromArray2` af_mean_all)

-- | Calculates weighted mean of all elements in an 'Array'

meanAllWeighted

:: AFType a

=> Array a

-- ^ Input 'Array'

-> Array a

-- ^ 'Array' of weights

-> (Double, Double)

-- ^ Weighted mean (real and imaginary part)

meanAllWeighted a b =

infoFromArray22 a b af_mean_all_weighted

-- | Calculates variance of all elements in an 'Array'

varAll

:: AFType a

=> Array a

-- ^ Input 'Array'

-> Bool

-- ^ Input 'Array'

-> (Double, Double)

-- ^ Variance (real and imaginary part)

varAll a (fromIntegral . fromEnum -> b) =

infoFromArray2 a $ \x y z ->

af_var_all x y z b

-- | Calculates weighted variance of all elements in an 'Array'

varAllWeighted

:: AFType a

=> Array a

-- ^ Input 'Array'

-> Array a

-- ^ 'Array' of weights

-> (Double, Double)

-- ^ Variance weighted result, (real and imaginary part)

varAllWeighted a b =

infoFromArray22 a b af_var_all_weighted

-- | Calculates standard deviation of all elements in an 'Array'

stdevAll

:: AFType a

=> Array a

-- ^ Input 'Array'

-> (Double, Double)

-- ^ Standard deviation result, (real and imaginary part)

stdevAll = (`infoFromArray2` af_stdev_all)

-- | Calculates median of all elements in an 'Array'

medianAll

:: (AFType a, Fractional a)

=> Array a

-- ^ Input 'Array'

-> (Double, Double)

-- ^ Median result, real and imaginary part

medianAll = (`infoFromArray2` af_median_all)

-- | This algorithm returns Pearson product-moment correlation coefficient.

corrCoef

:: AFType a

=> Array a

-- ^ First input 'Array'

-> Array a

-- ^ Second input 'Array'

-> (Double, Double)

-- ^ Correlation coefficient result, real and imaginary part

corrCoef a b =

infoFromArray22 a b af_corrcoef

-- | This function returns the top k values along a given dimension of the input array.

topk

:: AFType a

=> Array a

-- ^ First input 'Array', with at least /k/ elements along /dim/

-> Int

-- ^ The number of elements to be retrieved along the dim dimension

-> TopK

-- ^ If descending, the highest values are returned. Otherwise, the lowest values are returned

-> (Array a, Array a)

-- ^ Returns The values of the top k elements along the dim dimension

-- along with the indices of the top k elements along the dim dimension

topk a (fromIntegral -> x) (fromTopK -> f)

= a `op2p` (\b c d -> af_topk b c d x 0 f)