-- | Simple statistics code.
module ParkBench.Internal.Statistics
  ( Timed (..),
    Estimate (..),
    initialEstimate,
    updateEstimate,
    elapsed,
    stdev,
    variance,
    goodness,
    Roll (..),
  )
where

import ParkBench.Internal.Prelude

-- | A value that took a certan time to compute.
data Timed a = Timed
  { forall a. Timed a -> Rational
nanoseconds :: {-# UNPACK #-} !Rational,
    forall a. Timed a -> a
value :: !a
  }
  deriving stock (forall a b. a -> Timed b -> Timed a
forall a b. (a -> b) -> Timed a -> Timed b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> Timed b -> Timed a
$c<$ :: forall a b. a -> Timed b -> Timed a
fmap :: forall a b. (a -> b) -> Timed a -> Timed b
$cfmap :: forall a b. (a -> b) -> Timed a -> Timed b
Functor, Int -> Timed a -> ShowS
forall a. Show a => Int -> Timed a -> ShowS
forall a. Show a => [Timed a] -> ShowS
forall a. Show a => Timed a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Timed a] -> ShowS
$cshowList :: forall a. Show a => [Timed a] -> ShowS
show :: Timed a -> String
$cshow :: forall a. Show a => Timed a -> String
showsPrec :: Int -> Timed a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Timed a -> ShowS
Show)

instance Monoid a => Monoid (Timed a) where
  mempty :: Timed a
mempty = forall a. Rational -> a -> Timed a
Timed Rational
0 forall a. Monoid a => a
mempty
  mappend :: Timed a -> Timed a -> Timed a
mappend = forall a. Semigroup a => a -> a -> a
(<>)

instance Semigroup a => Semigroup (Timed a) where
  Timed Rational
n0 a
x0 <> :: Timed a -> Timed a -> Timed a
<> Timed Rational
n1 a
x1 =
    forall a. Rational -> a -> Timed a
Timed (Rational
n0 forall a. Num a => a -> a -> a
+ Rational
n1) (a
x0 forall a. Semigroup a => a -> a -> a
<> a
x1)

data Estimate a = Estimate
  { forall a. Estimate a -> Rational
kvariance :: {-# UNPACK #-} !Rational,
    forall a. Estimate a -> Timed a
mean :: {-# UNPACK #-} !(Timed a),
    forall a. Estimate a -> Word64
samples :: {-# UNPACK #-} !Word64
  }
  deriving stock (forall a b. a -> Estimate b -> Estimate a
forall a b. (a -> b) -> Estimate a -> Estimate b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> Estimate b -> Estimate a
$c<$ :: forall a b. a -> Estimate b -> Estimate a
fmap :: forall a b. (a -> b) -> Estimate a -> Estimate b
$cfmap :: forall a b. (a -> b) -> Estimate a -> Estimate b
Functor, Int -> Estimate a -> ShowS
forall a. Show a => Int -> Estimate a -> ShowS
forall a. Show a => [Estimate a] -> ShowS
forall a. Show a => Estimate a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Estimate a] -> ShowS
$cshowList :: forall a. Show a => [Estimate a] -> ShowS
show :: Estimate a -> String
$cshow :: forall a. Show a => Estimate a -> String
showsPrec :: Int -> Estimate a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Estimate a -> ShowS
Show)

-- | The total elapsed time of an estimate, in nanoseconds.
elapsed :: Estimate a -> Rational
elapsed :: forall a. Estimate a -> Rational
elapsed Estimate {Timed a
mean :: Timed a
$sel:mean:Estimate :: forall a. Estimate a -> Timed a
mean, Word64
samples :: Word64
$sel:samples:Estimate :: forall a. Estimate a -> Word64
samples} =
  Word64 -> Rational
w2r Word64
samples forall a. Num a => a -> a -> a
* forall a. Timed a -> Rational
nanoseconds Timed a
mean

-- | The standard deviation of an estimate.
stdev :: Estimate a -> Double
stdev :: forall a. Estimate a -> Double
stdev =
  forall a. Floating a => a -> a
sqrt forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Fractional a => Rational -> a
fromRational @Double) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Estimate a -> Rational
variance

-- | The variance of an estimate.
variance :: Estimate a -> Rational
variance :: forall a. Estimate a -> Rational
variance (Estimate Rational
kvariance Timed a
_ Word64
samples) =
  if Word64
samples forall a. Eq a => a -> a -> Bool
== Word64
1
    then Rational
0
    else Rational
kvariance forall a. Fractional a => a -> a -> a
/ Word64 -> Rational
w2r (Word64
samples forall a. Num a => a -> a -> a
- Word64
1)

-- | The "goodness" of an estimate, which is just how large its standard deviation is, relative to its mean.
--
-- Smaller is better, and the smallest possible value is 0.
goodness :: Estimate a -> Double
goodness :: forall a. Estimate a -> Double
goodness Estimate a
e =
  forall a. Estimate a -> Double
stdev Estimate a
e forall a. Fractional a => a -> a -> a
/ Rational -> Double
r2d (forall a. Timed a -> Rational
nanoseconds (forall a. Estimate a -> Timed a
mean Estimate a
e))

-- | @initialEstimate v@ creates an estimate per thing-that-took-time @v@ that was a run of 1 iteration.
initialEstimate :: Timed a -> Estimate a
initialEstimate :: forall a. Timed a -> Estimate a
initialEstimate Timed a
mean =
  Estimate
    { $sel:kvariance:Estimate :: Rational
kvariance = Rational
0,
      Timed a
mean :: Timed a
$sel:mean:Estimate :: Timed a
mean,
      $sel:samples:Estimate :: Word64
samples = Word64
1
    }

-- | @updateEstimate n v e@ updates estimate @e@ per thing-that-took-time @v@ that was a run of @n@ iterations.
updateEstimate :: Roll a => Word64 -> Timed a -> Estimate a -> Estimate a
updateEstimate :: forall a. Roll a => Word64 -> Timed a -> Estimate a -> Estimate a
updateEstimate Word64
n (Timed Rational
tn1 a
value1) (Estimate Rational
kvariance0 (Timed Rational
mean0 a
value0) Word64
samples0) =
  Estimate
    { $sel:kvariance:Estimate :: Rational
kvariance = Rational
kvariance0 forall a. Num a => a -> a -> a
+ Rational
nr forall a. Num a => a -> a -> a
* (Rational
t1 forall a. Num a => a -> a -> a
- Rational
mean0) forall a. Num a => a -> a -> a
* (Rational
t1 forall a. Num a => a -> a -> a
- Rational
mean1),
      $sel:mean:Estimate :: Timed a
mean = forall a. Rational -> a -> Timed a
Timed Rational
mean1 a
value',
      $sel:samples:Estimate :: Word64
samples = Word64
samples1
    }
  where
    mean1 :: Rational
mean1 = Rational -> Rational -> Rational
rollmean Rational
mean0 Rational
tn1
    samples1 :: Word64
samples1 = Word64
samples0 forall a. Num a => a -> a -> a
+ Word64
n
    t1 :: Rational
t1 = Rational
tn1 forall a. Fractional a => a -> a -> a
/ Rational
nr
    value' :: a
value' = forall a.
Roll a =>
(Rational -> Rational -> Rational) -> a -> a -> a
roll Rational -> Rational -> Rational
rollmean a
value0 a
value1
    rollmean :: Rational -> Rational -> Rational
rollmean Rational
u0 Rational
u1 = Rational
u0 forall a. Num a => a -> a -> a
+ ((Rational
u1 forall a. Num a => a -> a -> a
- Rational
nr forall a. Num a => a -> a -> a
* Rational
u0) forall a. Fractional a => a -> a -> a
/ Word64 -> Rational
w2r Word64
samples1)
    nr :: Rational
nr = Word64 -> Rational
w2r Word64
n

class Roll a where
  roll :: (Rational -> Rational -> Rational) -> a -> a -> a