{-# LANGUAGE
    MultiParamTypeClasses,
    FlexibleInstances, FlexibleContexts, UndecidableInstances,
    TemplateHaskell
  #-}

{-# OPTIONS_GHC -fno-warn-simplifiable-class-constraints #-}

module Data.Random.Distribution.Poisson where

import Data.Random.Internal.TH

import Data.Random.RVar
import Data.Random.Distribution
import Data.Random.Distribution.Uniform
import Data.Random.Distribution.Gamma
import Data.Random.Distribution.Binomial

import Control.Monad

-- from Knuth, with interpretation help from gsl sources
integralPoisson :: (Integral a, RealFloat b, Distribution StdUniform b, Distribution (Erlang a) b, Distribution (Binomial b) a) => b -> RVarT m a
integralPoisson :: b -> RVarT m a
integralPoisson = a -> b -> RVarT m a
forall a b (m :: * -> *).
(Integral a, RealFloat b, Distribution StdUniform b,
 Distribution (Erlang a) b, Distribution (Binomial b) a) =>
a -> b -> RVarT m a
psn 0
    where
        psn :: (Integral a, RealFloat b, Distribution StdUniform b, Distribution (Erlang a) b, Distribution (Binomial b) a) => a -> b -> RVarT m a
        psn :: a -> b -> RVarT m a
psn j :: a
j mu :: b
mu
            | b
mu b -> b -> Bool
forall a. Ord a => a -> a -> Bool
> 10   = do
                let m :: a
m = b -> a
forall a b. (RealFrac a, Integral b) => a -> b
floor (b
mu b -> b -> b
forall a. Num a => a -> a -> a
* (7b -> b -> b
forall a. Fractional a => a -> a -> a
/8))

                b
x <- a -> RVarT m b
forall a b (m :: * -> *).
Distribution (Erlang a) b =>
a -> RVarT m b
erlangT a
m
                if b
x b -> b -> Bool
forall a. Ord a => a -> a -> Bool
>= b
mu
                    then do
                        a
b <- a -> b -> RVarT m a
forall b a (m :: * -> *).
Distribution (Binomial b) a =>
a -> b -> RVarT m a
binomialT (a
m a -> a -> a
forall a. Num a => a -> a -> a
- 1) (b
mu b -> b -> b
forall a. Fractional a => a -> a -> a
/ b
x)
                        a -> RVarT m a
forall (m :: * -> *) a. Monad m => a -> m a
return (a
j a -> a -> a
forall a. Num a => a -> a -> a
+ a
b)
                    else a -> b -> RVarT m a
forall a b (m :: * -> *).
(Integral a, RealFloat b, Distribution StdUniform b,
 Distribution (Erlang a) b, Distribution (Binomial b) a) =>
a -> b -> RVarT m a
psn (a
j a -> a -> a
forall a. Num a => a -> a -> a
+ a
m) (b
mu b -> b -> b
forall a. Num a => a -> a -> a
- b
x)

            | Bool
otherwise = b -> a -> RVarT m a
forall b (m :: * -> *). Num b => b -> b -> RVarT m b
prod 1 a
j
                where
                    emu :: b
emu = b -> b
forall a. Floating a => a -> a
exp (-b
mu)

                    prod :: b -> b -> RVarT m b
prod p :: b
p k :: b
k = do
                        b
u <- RVarT m b
forall a (m :: * -> *). Distribution StdUniform a => RVarT m a
stdUniformT
                        if b
p b -> b -> b
forall a. Num a => a -> a -> a
* b
u b -> b -> Bool
forall a. Ord a => a -> a -> Bool
> b
emu
                            then b -> b -> RVarT m b
prod (b
p b -> b -> b
forall a. Num a => a -> a -> a
* b
u) (b
k b -> b -> b
forall a. Num a => a -> a -> a
+ 1)
                            else b -> RVarT m b
forall (m :: * -> *) a. Monad m => a -> m a
return b
k

integralPoissonCDF :: (Integral a, Real b) => b -> a -> Double
integralPoissonCDF :: b -> a -> Double
integralPoissonCDF mu :: b
mu k :: a
k = Double -> Double
forall a. Floating a => a -> a
exp (Double -> Double
forall a. Num a => a -> a
negate Double
lambda) Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Double] -> Double
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum
    [ Double -> Double
forall a. Floating a => a -> a
exp (a -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
i Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log Double
lambda Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
i_fac_ln)
    | (i :: a
i, i_fac_ln :: Double
i_fac_ln) <- [a] -> [Double] -> [(a, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [0..a
k] ((Double -> Double -> Double) -> Double -> [Double] -> [Double]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Double -> Double -> Double
forall a. Num a => a -> a -> a
(+) 0 ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
forall a. Floating a => a -> a
log [1..]))
    ]

    where lambda :: Double
lambda = b -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac b
mu

-- | The probability of getting exactly k successes is
-- given by the probability mass function:
--
-- \[
-- f(k;\lambda) = \Pr(X = k) = \frac{\lambda^k e^{-\lambda}}{k!}
-- \]
--
-- Note that in `integralPoissonPDF` the parameter of the mass
-- function are given first and the range of the random variable
-- distributed according to the Poisson distribution is given
-- last. That is, \(f(2;0.5)\) is calculated by @integralPoissonPDF 0.5 2@.
integralPoissonPDF :: (Integral a, Real b) => b -> a -> Double
integralPoissonPDF :: b -> a -> Double
integralPoissonPDF mu :: b
mu k :: a
k = Double -> Double
forall a. Floating a => a -> a
exp (Double -> Double
forall a. Num a => a -> a
negate Double
lambda) Double -> Double -> Double
forall a. Num a => a -> a -> a
*
                          Double -> Double
forall a. Floating a => a -> a
exp (a -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
k Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log Double
lambda Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
k_fac_ln)
  where
    k_fac_ln :: Double
k_fac_ln = (Double -> Double -> Double) -> Double -> [Double] -> Double
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Double -> Double -> Double
forall a. Num a => a -> a -> a
(+) 0 ((a -> Double) -> [a] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double
forall a. Floating a => a -> a
log (Double -> Double) -> (a -> Double) -> a -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral) [1..a
k])
    lambda :: Double
lambda   = b -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac b
mu

fractionalPoisson :: (Num a, Distribution (Poisson b) Integer) => b -> RVarT m a
fractionalPoisson :: b -> RVarT m a
fractionalPoisson mu :: b
mu = (Integer -> a) -> RVarT m Integer -> RVarT m a
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM Integer -> a
forall a. Num a => Integer -> a
fromInteger (b -> RVarT m Integer
forall b a (m :: * -> *).
Distribution (Poisson b) a =>
b -> RVarT m a
poissonT b
mu)

fractionalPoissonCDF :: (CDF (Poisson b) Integer, RealFrac a) => b -> a -> Double
fractionalPoissonCDF :: b -> a -> Double
fractionalPoissonCDF mu :: b
mu k :: a
k = Poisson b Integer -> Integer -> Double
forall (d :: * -> *) t. CDF d t => d t -> t -> Double
cdf (b -> Poisson b Integer
forall b a. b -> Poisson b a
Poisson b
mu) (a -> Integer
forall a b. (RealFrac a, Integral b) => a -> b
floor a
k :: Integer)

fractionalPoissonPDF :: (PDF (Poisson b) Integer, RealFrac a) => b -> a -> Double
fractionalPoissonPDF :: b -> a -> Double
fractionalPoissonPDF mu :: b
mu k :: a
k = Poisson b Integer -> Integer -> Double
forall (d :: * -> *) t. PDF d t => d t -> t -> Double
pdf (b -> Poisson b Integer
forall b a. b -> Poisson b a
Poisson b
mu) (a -> Integer
forall a b. (RealFrac a, Integral b) => a -> b
floor a
k :: Integer)

poisson :: (Distribution (Poisson b) a) => b -> RVar a
poisson :: b -> RVar a
poisson mu :: b
mu = Poisson b a -> RVar a
forall (d :: * -> *) t. Distribution d t => d t -> RVar t
rvar (b -> Poisson b a
forall b a. b -> Poisson b a
Poisson b
mu)

poissonT :: (Distribution (Poisson b) a) => b -> RVarT m a
poissonT :: b -> RVarT m a
poissonT mu :: b
mu = Poisson b a -> RVarT m a
forall (d :: * -> *) t (n :: * -> *).
Distribution d t =>
d t -> RVarT n t
rvarT (b -> Poisson b a
forall b a. b -> Poisson b a
Poisson b
mu)

newtype Poisson b a = Poisson b

$( replicateInstances ''Int integralTypes [d|
        instance ( RealFloat b
                 , Distribution StdUniform   b
                 , Distribution (Erlang Int) b
                 , Distribution (Binomial b) Int
                 ) => Distribution (Poisson b) Int where
            rvarT (Poisson mu) = integralPoisson mu
        instance (Real b, Distribution (Poisson b) Int) => CDF (Poisson b) Int where
            cdf  (Poisson mu) = integralPoissonCDF mu
        instance (Real b, Distribution (Poisson b) Int) => PDF (Poisson b) Int where
            pdf  (Poisson mu) = integralPoissonPDF mu
    |] )

$( replicateInstances ''Float realFloatTypes [d|
        instance (Distribution (Poisson b) Integer) => Distribution (Poisson b) Float where
            rvarT (Poisson mu) = fractionalPoisson mu
        instance (CDF (Poisson b) Integer) => CDF (Poisson b) Float where
            cdf  (Poisson mu) = fractionalPoissonCDF mu
        instance (PDF (Poisson b) Integer) => PDF (Poisson b) Float where
            pdf  (Poisson mu) = fractionalPoissonPDF mu
    |])