-- SPDX-License-Identifier: MPL-2.0

module HMath.Statistics (
    mean,
    variance,
    standardDeviation,
    median,
) where

import Data.List (sort)

{- | Arithmetic mean of a finite nonempty list: sum divided by length. Empty input raises a
pure error; no non-finite-value validation is performed. Example: mean [2,4] == 3.
-}
mean :: (Fractional a) => [a] -> a
mean :: forall a. Fractional a => [a] -> a
mean [] = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"mean: empty input"
mean [a]
xs = [a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [a]
xs a -> a -> a
forall a. Fractional a => a -> a -> a
/ Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs)

{- | Population variance of a finite nonempty list, dividing squared deviations by n (not
n-1). Empty input raises a pure error. Example: variance [1,3] == 1.
-}
variance :: (Fractional a) => [a] -> a
variance :: forall a. Fractional a => [a] -> a
variance [] = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"variance: empty input"
variance [a]
xs =
    let m :: a
m = [a] -> a
forall a. Fractional a => [a] -> a
mean [a]
xs
        squared :: [a]
squared = (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (\a
x -> (a
x a -> a -> a
forall a. Num a => a -> a -> a
- a
m) a -> a -> a
forall a. Num a => a -> a -> a
* (a
x a -> a -> a
forall a. Num a => a -> a -> a
- a
m)) [a]
xs
     in [a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [a]
squared a -> a -> a
forall a. Fractional a => a -> a -> a
/ Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs)

{- | Population standard deviation, sqrt of variance. Requires a finite nonempty list; empty
input raises the same error as variance. Example: standardDeviation [1,3] == 1.
-}
standardDeviation :: (Floating a) => [a] -> a
standardDeviation :: forall a. Floating a => [a] -> a
standardDeviation = a -> a
forall a. Floating a => a -> a
sqrt (a -> a) -> ([a] -> a) -> [a] -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> a
forall a. Fractional a => [a] -> a
variance

{- | Sort a finite nonempty list and return its middle value, averaging the central pair for
even length. Empty input raises a pure error. Example: median [4,1,3,2] == 2.5.
-}
median :: (Ord a, Fractional a) => [a] -> a
median :: forall a. (Ord a, Fractional a) => [a] -> a
median [] = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"median: empty input"
median [a]
xs
    | Int -> Bool
forall a. Integral a => a -> Bool
odd Int
n = [a]
sorted [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
mid
    | Bool
otherwise = ([a]
sorted [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
mid Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) a -> a -> a
forall a. Num a => a -> a -> a
+ [a]
sorted [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
mid) a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
2
  where
    sorted :: [a]
sorted = [a] -> [a]
forall a. Ord a => [a] -> [a]
sort [a]
xs
    n :: Int
n = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs
    mid :: Int
mid = Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2