-- |
-- Module      :  Numeric.Stats
-- Copyright   :  (c) OleksandrZhabenko 2020
-- License     :  MIT
-- Stability   :  Experimental
-- Maintainer  :  olexandr543@yahoo.com
--
-- A very basic descriptive statistics. Functions use a tail recursion approach to compute the values and are strict by an accumulator.

{-# LANGUAGE BangPatterns, MagicHash #-}

module Numeric.Stats where

import GHC.Exts
import GHC.Prim

-- | A tail-recursive realization for the statistic mean. 
mean2 :: RealFrac a => [a] -> a -> a -> a -> a
mean2 (x:xs) !s1 !l1 _ = mean2 xs (s1 + x) (l1 + 1) ((s1 + x) / (l1 + 1))
mean2 _ _ _ m = m

-- | Similar to 'mean2', but uses GHC unlifted types from @ghc-prim@ package. 
mean2F :: [Float] -> Float# -> Float# -> Float# -> Float
mean2F ((F# !x):xs) !s1 !l1 _ = mean2F xs (plusFloat# s1 x) (plusFloat# l1 1.0#) (divideFloat# (plusFloat# s1 x) (plusFloat# l1 1.0#))
mean2F _ _ _ m = F# m

-- | One-pass and tail-recursive realization for the pair of the mean and dispersion. Is vulnerable to the floating-point cancellation errors. 
meanWithDispersion :: (RealFrac a, Floating a) => [a] -> a -> a -> a -> a -> a -> (a,a)
meanWithDispersion (!x:xs) !s1 !s2 !l1 m1 d = meanWithDispersion xs (s1 + x) (s2 + x*x) (l1 + 1) (m0 s1 l1 x) (m0 s2 l1 (x*x) - (m0 s1 l1 x)**2)
  where m0 !s3 !l2 !x = (s3 + x) / (l2 + 1)
meanWithDispersion _ _ _ _ !m !d = (m,d)

-- | A tail-recursive realization for the statistic mean. A variant of the 'mean2' function. 
mean :: RealFrac a => [a] -> a
mean xs = mean2 xs 0.0 0.0 0.0
{-# INLINE mean #-}

-- | Uses 'mean2F' inside. 
meanF :: [Float] -> Float
meanF xs = mean2F xs 0.0# 0.0# 0.0#
{-# INLINE meanF #-}

meanWithDisp :: (RealFrac a, Floating a) => [a] -> (a,a)
meanWithDisp xs = meanWithDispersion xs 0.0 0.0 0.0 0.0 0.0
{-# INLINE meanWithDisp #-}