module Data.TDigest.Postprocess (
    -- * Histogram
    I.HasHistogram (..),
    I.HistBin (..),
    -- * Quantiles
    median,
    quantile,
    -- * Mean & variance
    --
    -- | As we have "full" histogram, we can calculate other statistical
    -- variables.
    mean,
    variance,
    stddev,
    -- * CDF
    cdf,
    icdf,
    -- * Affine
    I.Affine (..)
    ) where

import qualified Data.List.NonEmpty as NE
import           Prelude ()
import           Prelude.Compat

import qualified Data.TDigest.Postprocess.Internal as I

-- | Median, i.e. @'quantile' 0.5@.
median :: I.HasHistogram a f => a -> f Double
median :: forall a (f :: * -> *). HasHistogram a f => a -> f Double
median = forall a (f :: * -> *). HasHistogram a f => Double -> a -> f Double
quantile Double
0.5

-- | Calculate quantile of a specific value.
quantile :: I.HasHistogram a f => Double -> a -> f Double
quantile :: forall a (f :: * -> *). HasHistogram a f => Double -> a -> f Double
quantile Double
q a
x = Double -> Double -> NonEmpty HistBin -> Double
I.quantile Double
q (forall a (f :: * -> *). HasHistogram a f => a -> Double
I.totalWeight a
x) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a (f :: * -> *).
HasHistogram a f =>
a -> f (NonEmpty HistBin)
I.histogram a
x

-- | Mean.
--
-- >>> mean (Tree.tdigest [1..100] :: Tree.TDigest 10)
-- Just 50.5
--
-- /Note:/ if you only need the mean, calculate it directly.
--
mean :: I.HasHistogram a f => a -> f Double
mean :: forall a (f :: * -> *). HasHistogram a f => a -> f Double
mean a
x = NonEmpty HistBin -> Double
I.mean forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a (f :: * -> *).
HasHistogram a f =>
a -> f (NonEmpty HistBin)
I.histogram a
x

-- | Variance.
--
variance :: I.HasHistogram a f => a -> f Double
variance :: forall a (f :: * -> *). HasHistogram a f => a -> f Double
variance a
x = NonEmpty HistBin -> Double
I.variance forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a (f :: * -> *).
HasHistogram a f =>
a -> f (NonEmpty HistBin)
I.histogram a
x

-- | Standard deviation, square root of variance.
stddev :: I.HasHistogram a f => a -> f Double
stddev :: forall a (f :: * -> *). HasHistogram a f => a -> f Double
stddev = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. Floating a => a -> a
sqrt forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a (f :: * -> *). HasHistogram a f => a -> f Double
variance

-- | Cumulative distribution function.
--
-- /Note:/ if this is the only thing you need, it's more efficient to count
-- this directly.
cdf :: I.HasHistogram a f => Double -> a -> Double
cdf :: forall a (f :: * -> *). HasHistogram a f => Double -> a -> Double
cdf Double
q a
x = forall (t :: * -> *) b a. Affine t => b -> (a -> b) -> t a -> b
I.affine Double
1 (Double -> Double -> [HistBin] -> Double
I.cdf Double
q (forall a (f :: * -> *). HasHistogram a f => a -> Double
I.totalWeight a
x) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. NonEmpty a -> [a]
NE.toList) forall a b. (a -> b) -> a -> b
$ forall a (f :: * -> *).
HasHistogram a f =>
a -> f (NonEmpty HistBin)
I.histogram a
x

-- | An alias for 'quantile'.
icdf :: I.HasHistogram a f => Double -> a -> f Double
icdf :: forall a (f :: * -> *). HasHistogram a f => Double -> a -> f Double
icdf = forall a (f :: * -> *). HasHistogram a f => Double -> a -> f Double
quantile

-- $setup
-- >>> :set -XDataKinds
-- >>> import qualified Data.TDigest.Tree as Tree