{-# LANGUAGE GeneralizedNewtypeDeriving #-}

{-|
Copyright   : Predictable Network Solutions Ltd., 2020-2024
License     : BSD-3-Clause
Description : Probability measures on the number line.
-}
module Numeric.Measure.Probability
    ( -- * Type
      Prob
    , dirac
    , uniform
    , distribution
    , fromDistribution
    , measure
    , fromMeasure
    , unsafeFromMeasure

      -- * Observations
    , support
    , expectation
    , moments

      -- * Operations, numerical
    , choice
    , translate
    , convolve
    ) where

import Control.DeepSeq
    ( NFData
    )
import Numeric.Function.Piecewise
    ( Piecewise
    )
import Numeric.Measure.Finite.Mixed
    ( Measure
    )
import Numeric.Polynomial.Simple
    ( Poly
    )
import Numeric.Probability.Moments
    ( Moments (..)
    , fromExpectedPowers
    )

import qualified Numeric.Measure.Finite.Mixed as M
import qualified Numeric.Polynomial.Simple as Poly

{-----------------------------------------------------------------------------
    Type
------------------------------------------------------------------------------}
-- | A
-- [probability measure](https://en.wikipedia.org/wiki/Probability_measure)
-- on the number line.
--
-- A probability measure is a 'M.Measure' whose 'M.total' is @1@.
newtype Prob a = Prob (Measure a)
    -- INVARIANT: 'M.isPositive' equals 'True'.
    -- INVARIANT: 'M.total' equals 1
    deriving (Int -> Prob a -> ShowS
[Prob a] -> ShowS
Prob a -> String
(Int -> Prob a -> ShowS)
-> (Prob a -> String) -> ([Prob a] -> ShowS) -> Show (Prob a)
forall a. Show a => Int -> Prob a -> ShowS
forall a. Show a => [Prob a] -> ShowS
forall a. Show a => Prob a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Prob a -> ShowS
showsPrec :: Int -> Prob a -> ShowS
$cshow :: forall a. Show a => Prob a -> String
show :: Prob a -> String
$cshowList :: forall a. Show a => [Prob a] -> ShowS
showList :: [Prob a] -> ShowS
Show, Prob a -> ()
(Prob a -> ()) -> NFData (Prob a)
forall a. NFData a => Prob a -> ()
forall a. (a -> ()) -> NFData a
$crnf :: forall a. NFData a => Prob a -> ()
rnf :: Prob a -> ()
NFData)

-- | View the probability measure as a 'M.Measure'.
measure :: (Ord a, Num a) => Prob a -> Measure a
measure :: forall a. (Ord a, Num a) => Prob a -> Measure a
measure (Prob Measure a
m) = Measure a
m

-- | View a 'M.Measure' as a probability distribution.
--
-- The measure @m@ must be positive, with total weight @1@, that is
--
-- > isPositive m == True
-- > total m == 1
--
-- These preconditions are checked and the function returns 'Nothing'
-- if they fail. 
fromMeasure :: (Ord a, Num a, Fractional a) => Measure a -> Maybe (Prob a)
fromMeasure :: forall a.
(Ord a, Num a, Fractional a) =>
Measure a -> Maybe (Prob a)
fromMeasure Measure a
m
    | Measure a -> Bool
forall a. (Ord a, Num a, Fractional a) => Measure a -> Bool
M.isPositive Measure a
m Bool -> Bool -> Bool
&& Measure a -> a
forall a. (Ord a, Num a) => Measure a -> a
M.total Measure a
m a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
1 = Prob a -> Maybe (Prob a)
forall a. a -> Maybe a
Just (Prob a -> Maybe (Prob a)) -> Prob a -> Maybe (Prob a)
forall a b. (a -> b) -> a -> b
$ Measure a -> Prob a
forall a. Measure a -> Prob a
Prob Measure a
m
    | Bool
otherwise = Maybe (Prob a)
forall a. Maybe a
Nothing

-- | View a 'M.Measure' as a probability distribution.
--
-- Variant of 'fromMeasure' where /the precondition are not checked!/
unsafeFromMeasure :: Measure a -> Prob a
unsafeFromMeasure :: forall a. Measure a -> Prob a
unsafeFromMeasure = Measure a -> Prob a
forall a. Measure a -> Prob a
Prob

-- | @eval (distribution m) x@ is the probability of picking a number @<= x@.
--
-- This is known as the
-- [cumulative distribution function
-- ](https://en.wikipedia.org/wiki/Cumulative_distribution_function).
distribution :: (Ord a, Num a) => Prob a -> Piecewise (Poly a)
distribution :: forall a. (Ord a, Num a) => Prob a -> Piecewise (Poly a)
distribution (Prob Measure a
m) = Measure a -> Piecewise (Poly a)
forall a. (Ord a, Num a) => Measure a -> Piecewise (Poly a)
M.distribution Measure a
m

-- | Construct a probability distribution from its
-- [cumulative distribution function
-- ](https://en.wikipedia.org/wiki/Cumulative_distribution_function).
--
-- Return 'Nothing' if
-- * the cumulative distribution function is not monotonicall increasing
-- * the last piece of the piecewise function is not a constant
--   equal to @1@.
fromDistribution
    :: (Ord a, Num a, Fractional a)
    => Piecewise (Poly a) -> Maybe (Prob a)
fromDistribution :: forall a.
(Ord a, Num a, Fractional a) =>
Piecewise (Poly a) -> Maybe (Prob a)
fromDistribution Piecewise (Poly a)
pieces
    | Just Measure a
m <- Piecewise (Poly a) -> Maybe (Measure a)
forall a. (Ord a, Num a) => Piecewise (Poly a) -> Maybe (Measure a)
M.fromDistribution Piecewise (Poly a)
pieces = Measure a -> Maybe (Prob a)
forall a.
(Ord a, Num a, Fractional a) =>
Measure a -> Maybe (Prob a)
fromMeasure Measure a
m
    | Bool
otherwise = Maybe (Prob a)
forall a. Maybe a
Nothing

-- | Two probability measures are equal if they have the same cumulative
-- distribution functions.
--
-- > px == py
-- >   implies
-- >   forall t. eval (distribution px) t = eval (distribution py) t
instance (Ord a, Num a) => Eq (Prob a) where
    Prob Measure a
mx == :: Prob a -> Prob a -> Bool
== Prob Measure a
my = Measure a
mx Measure a -> Measure a -> Bool
forall a. Eq a => a -> a -> Bool
== Measure a
my

{-----------------------------------------------------------------------------
    Construction
------------------------------------------------------------------------------}
-- | A
-- [Dirac measure](https://en.wikipedia.org/wiki/Dirac_measure)
-- at the given point @x@.
--
-- @dirac x@ is the probability distribution where @x@ occurs with certainty.
dirac :: (Ord a, Num a) => a -> Prob a
dirac :: forall a. (Ord a, Num a) => a -> Prob a
dirac = Measure a -> Prob a
forall a. Measure a -> Prob a
Prob (Measure a -> Prob a) -> (a -> Measure a) -> a -> Prob a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Measure a
forall a. (Ord a, Num a) => a -> Measure a
M.dirac

-- | The uniform probability distribution on the interval \( [x,y) \).
uniform :: (Ord a, Num a, Fractional a) => a -> a -> Prob a
uniform :: forall a. (Ord a, Num a, Fractional a) => a -> a -> Prob a
uniform a
x a
y = Measure a -> Prob a
forall a. Measure a -> Prob a
Prob (Measure a -> Prob a) -> Measure a -> Prob a
forall a b. (a -> b) -> a -> b
$ a -> a -> Measure a
forall a. (Ord a, Num a, Fractional a) => a -> a -> Measure a
M.uniform a
x a
y

{-----------------------------------------------------------------------------
    Construction
------------------------------------------------------------------------------}
-- | The 'support' is the smallest closed, contiguous interval \( [x,y] \)
-- outside of which the probability is zero.
--
-- Returns 'Nothing' if the interval is empty.
support :: (Ord a, Num a) => Prob a -> Maybe (a, a)
support :: forall a. (Ord a, Num a) => Prob a -> Maybe (a, a)
support (Prob Measure a
m) = Measure a -> Maybe (a, a)
forall a. (Ord a, Num a) => Measure a -> Maybe (a, a)
M.support Measure a
m

-- | Compute the
-- [expected value](https://en.wikipedia.org/wiki/Expected_value)
-- of a polynomial @f@ with respect to the given probability distribution,
-- \( E[f(X)] \).
expectation :: (Ord a, Num a, Fractional a) => Poly a -> Prob a -> a
expectation :: forall a. (Ord a, Num a, Fractional a) => Poly a -> Prob a -> a
expectation Poly a
f (Prob Measure a
m) = Poly a -> Measure a -> a
forall a. (Ord a, Num a, Fractional a) => Poly a -> Measure a -> a
M.integrate Poly a
f Measure a
m

-- | Compute the first four
-- commonly used moments of a probability distribution.
moments :: (Ord a, Num a, Fractional a) => Prob a -> Moments a
moments :: forall a. (Ord a, Num a, Fractional a) => Prob a -> Moments a
moments Prob a
m =
    (a, a, a, a) -> Moments a
forall a. (Ord a, Num a, Fractional a) => (a, a, a, a) -> Moments a
fromExpectedPowers (Int -> a
ex Int
1, Int -> a
ex Int
2, Int -> a
ex Int
3, Int -> a
ex Int
4)
  where
    ex :: Int -> a
ex Int
n = Poly a -> Prob a -> a
forall a. (Ord a, Num a, Fractional a) => Poly a -> Prob a -> a
expectation (Int -> a -> Poly a
forall a. (Eq a, Num a) => Int -> a -> Poly a
Poly.monomial Int
n a
1) Prob a
m

{-----------------------------------------------------------------------------
    Operations
------------------------------------------------------------------------------}
-- | Left-biased random choice.
--
-- @choice p@ is a probability distribution where
-- events from the left argument are chosen with probablity @p@
-- and events from the right argument are chosen with probability @(1-p)@.
--
-- > eval (distribution (choice p mx my)) z
-- >    = p * eval (distribution mx) z + (1-p) * eval (distribution my) z
choice :: (Ord a, Num a, Fractional a) => a -> Prob a -> Prob a -> Prob a
choice :: forall a.
(Ord a, Num a, Fractional a) =>
a -> Prob a -> Prob a -> Prob a
choice a
p (Prob Measure a
mx) (Prob Measure a
my) = Measure a -> Prob a
forall a. Measure a -> Prob a
Prob (Measure a -> Prob a) -> Measure a -> Prob a
forall a b. (a -> b) -> a -> b
$
    Measure a -> Measure a -> Measure a
forall a. (Ord a, Num a) => Measure a -> Measure a -> Measure a
M.add (a -> Measure a -> Measure a
forall a. (Ord a, Num a) => a -> Measure a -> Measure a
M.scale a
p Measure a
mx) (a -> Measure a -> Measure a
forall a. (Ord a, Num a) => a -> Measure a -> Measure a
M.scale (a
1 a -> a -> a
forall a. Num a => a -> a -> a
- a
p) Measure a
my)

-- | Translate a probability distribution along the number line.
--
-- > eval (distribution (translate y m)) x
-- >    = eval (distribution m) (x - y)
translate :: (Ord a, Num a, Fractional a) => a -> Prob a -> Prob a
translate :: forall a. (Ord a, Num a, Fractional a) => a -> Prob a -> Prob a
translate a
y (Prob Measure a
m) = Measure a -> Prob a
forall a. Measure a -> Prob a
Prob (Measure a -> Prob a) -> Measure a -> Prob a
forall a b. (a -> b) -> a -> b
$ a -> Measure a -> Measure a
forall a.
(Ord a, Num a, Fractional a) =>
a -> Measure a -> Measure a
M.translate a
y Measure a
m

-- | Additive convolution of two probability measures.
--
-- Properties:
--
-- > convolve (dirac x) (dirac y) = dirac (x + y)
-- >
-- > convolve mx my               =  convolve my mx
-- > translate z (convolve mx my) =  convolve (translate z mx) my
convolve
    :: (Ord a, Num a, Fractional a)
    => Prob a -> Prob a -> Prob a
convolve :: forall a.
(Ord a, Num a, Fractional a) =>
Prob a -> Prob a -> Prob a
convolve (Prob Measure a
mx) (Prob Measure a
my) = Measure a -> Prob a
forall a. Measure a -> Prob a
Prob (Measure a -> Prob a) -> Measure a -> Prob a
forall a b. (a -> b) -> a -> b
$ Measure a -> Measure a -> Measure a
forall a.
(Ord a, Num a, Fractional a) =>
Measure a -> Measure a -> Measure a
M.convolve Measure a
mx Measure a
my