{-# LANGUAGE NamedFieldPuns #-}

{-|
Copyright   : Predictable Network Solutions Ltd., 2020-2024
License     : BSD-3-Clause
Description : Moments of probability distributions.
-}
module Numeric.Probability.Moments
    ( Moments (..)
    , fromExpectedPowers
    ) where

{-----------------------------------------------------------------------------
    Test
------------------------------------------------------------------------------}

-- | The first four commonly used moments of a probability distribution.
data Moments a = Moments
    { forall a. Moments a -> a
mean :: a
    -- ^ [Mean or Expected Value](https://en.wikipedia.org/wiki/Expected_value)
    -- \( \mu \).
    -- Defined as \( \mu = E[X] \).
    , forall a. Moments a -> a
variance :: a
    -- ^ [Variance](https://en.wikipedia.org/wiki/Variance) \( \sigma^2 \).
    -- Defined as \( \sigma^2 = E[(X - \mu)^2] \).
    -- Equal to \( \sigma^2 = E[X^2] - \mu^2 \).
    , forall a. Moments a -> a
skewness :: a
    -- ^ [Skewness](https://en.wikipedia.org/wiki/Skewness) \( \gamma_1 \).
    -- Defined as
    -- \( \gamma_1 = E\left[\left(\frac{(X - \mu)}{\sigma}\right)^3 \right] \).
    , forall a. Moments a -> a
kurtosis :: a
    -- ^ [Kurtosis](https://en.wikipedia.org/wiki/Kurtosis) \( \kappa \).
    -- Defined as
    --  \( \kappa = E\left[\left(\frac{(X - \mu)}{\sigma}\right)^4 \right] \).
    --
    -- The kurtosis is bounded below: \( \kappa \geq \gamma_1^2 + 1 \).
    }
    deriving (Moments a -> Moments a -> Bool
(Moments a -> Moments a -> Bool)
-> (Moments a -> Moments a -> Bool) -> Eq (Moments a)
forall a. Eq a => Moments a -> Moments a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Moments a -> Moments a -> Bool
== :: Moments a -> Moments a -> Bool
$c/= :: forall a. Eq a => Moments a -> Moments a -> Bool
/= :: Moments a -> Moments a -> Bool
Eq, Int -> Moments a -> ShowS
[Moments a] -> ShowS
Moments a -> String
(Int -> Moments a -> ShowS)
-> (Moments a -> String)
-> ([Moments a] -> ShowS)
-> Show (Moments a)
forall a. Show a => Int -> Moments a -> ShowS
forall a. Show a => [Moments a] -> ShowS
forall a. Show a => Moments a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Moments a -> ShowS
showsPrec :: Int -> Moments a -> ShowS
$cshow :: forall a. Show a => Moments a -> String
show :: Moments a -> String
$cshowList :: forall a. Show a => [Moments a] -> ShowS
showList :: [Moments a] -> ShowS
Show)

-- | Compute the 'Moments' of a probability distribution given
-- the expectation values of the first four powers \( m_k = E[X^k] \).
--
-- > fromExpectedPowers (m1,m2,m3,m4)
fromExpectedPowers
    :: (Ord a, Num a, Fractional a)
    => (a, a, a, a) -> Moments a
fromExpectedPowers :: forall a. (Ord a, Num a, Fractional a) => (a, a, a, a) -> Moments a
fromExpectedPowers (a
mean, a
m2, a
m3, a
m4)
    | a
variance a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 =
        Moments{a
mean :: a
mean :: a
mean, a
variance :: a
variance :: a
variance, skewness :: a
skewness = a
0, kurtosis :: a
kurtosis = a
1}
    | Bool
otherwise =
        Moments{a
mean :: a
mean :: a
mean, a
variance :: a
variance :: a
variance, a
skewness :: a
skewness :: a
skewness, a
kurtosis :: a
kurtosis :: a
kurtosis}
  where
    meanSq :: a
meanSq = a
mean a -> a -> a
forall a. Num a => a -> a -> a
* a
mean

    variance :: a
variance = a
m2 a -> a -> a
forall a. Num a => a -> a -> a
- a
meanSq
    sigma :: a
sigma = a -> a
forall a. (Ord a, Num a, Fractional a) => a -> a
squareRoot a
variance

    skewness :: a
skewness =
        (a
m3 a -> a -> a
forall a. Num a => a -> a -> a
- a
3 a -> a -> a
forall a. Num a => a -> a -> a
* a
mean a -> a -> a
forall a. Num a => a -> a -> a
* a
variance a -> a -> a
forall a. Num a => a -> a -> a
- a
mean a -> a -> a
forall a. Num a => a -> a -> a
* a
meanSq
        ) a -> a -> a
forall a. Fractional a => a -> a -> a
/ (a
sigma a -> a -> a
forall a. Num a => a -> a -> a
* a
variance)

    kurtosis :: a
kurtosis =
        (a
m4
            a -> a -> a
forall a. Num a => a -> a -> a
- a
4 a -> a -> a
forall a. Num a => a -> a -> a
* a
mean a -> a -> a
forall a. Num a => a -> a -> a
* a
skewness a -> a -> a
forall a. Num a => a -> a -> a
* a
sigma a -> a -> a
forall a. Num a => a -> a -> a
* a
variance
            a -> a -> a
forall a. Num a => a -> a -> a
- a
6 a -> a -> a
forall a. Num a => a -> a -> a
* a
meanSq a -> a -> a
forall a. Num a => a -> a -> a
* a
variance
            a -> a -> a
forall a. Num a => a -> a -> a
- a
meanSq a -> a -> a
forall a. Num a => a -> a -> a
* a
meanSq
        ) a -> a -> a
forall a. Fractional a => a -> a -> a
/ (a
variance a -> a -> a
forall a. Num a => a -> a -> a
* a
variance)

-- | Helper function to approximate the square root.
-- Precision: 1e-4 of the given value.
--
-- Uses Heron's iterative method.
squareRoot :: (Ord a, Num a, Fractional a) => a -> a
squareRoot :: forall a. (Ord a, Num a, Fractional a) => a -> a
squareRoot a
x
    | a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0 = String -> a
forall a. HasCallStack => String -> a
error String
"Negative square root input"
    | a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 = a
0
    | Bool
otherwise = a -> a
goRoot a
x0
  where
    precision :: a
precision = a
x a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
10000
    x0 :: a
x0 = a
xa -> a -> a
forall a. Fractional a => a -> a -> a
/a
2 -- initial guess
    goRoot :: a -> a
goRoot a
xi
        | a -> a
forall a. Num a => a -> a
abs (a
x a -> a -> a
forall a. Num a => a -> a -> a
- a
xi a -> a -> a
forall a. Num a => a -> a -> a
* a
xi) a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
precision = a
xi
        | Bool
otherwise = a -> a
goRoot ((a
xi a -> a -> a
forall a. Num a => a -> a -> a
+ a
x a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
xi)a -> a -> a
forall a. Fractional a => a -> a -> a
/a
2)

{-sqRoot :: a -> a
sqRoot x = 
    let
        y :: Double
        y = toRational x
    in fromRational . toRational . sqrt y
-}