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

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

module Data.Random.Distribution.Poisson where

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

import Data.Int
import Data.Word

-- 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 :: forall a b (m :: * -> *).
(Integral a, RealFloat b, Distribution StdUniform b,
 Distribution (Erlang a) b, Distribution (Binomial b) a) =>
b -> RVarT m a
integralPoisson = 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
0
    where
        psn :: (Integral a, RealFloat b, Distribution StdUniform b, Distribution (Erlang a) b, Distribution (Binomial b) a) => a -> b -> RVarT m a
        psn :: 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 b
mu
            | b
mu forall a. Ord a => a -> a -> Bool
> b
10   = do
                let m :: a
m = forall a b. (RealFrac a, Integral b) => a -> b
floor (b
mu forall a. Num a => a -> a -> a
* (b
7forall a. Fractional a => a -> a -> a
/b
8))

                b
x <- forall a b (m :: * -> *).
Distribution (Erlang a) b =>
a -> RVarT m b
erlangT a
m
                if b
x forall a. Ord a => a -> a -> Bool
>= b
mu
                    then do
                        a
b <- forall b a (m :: * -> *).
Distribution (Binomial b) a =>
a -> b -> RVarT m a
binomialT (a
m forall a. Num a => a -> a -> a
- a
1) (b
mu forall a. Fractional a => a -> a -> a
/ b
x)
                        forall (m :: * -> *) a. Monad m => a -> m a
return (a
j forall a. Num a => a -> a -> a
+ a
b)
                    else 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 forall a. Num a => a -> a -> a
+ a
m) (b
mu forall a. Num a => a -> a -> a
- b
x)

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

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

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

    where lambda :: Double
lambda = 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 :: forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonPDF b
mu a
k = forall a. Floating a => a -> a
exp (forall a. Num a => a -> a
negate Double
lambda) forall a. Num a => a -> a -> a
*
                          forall a. Floating a => a -> a
exp (forall a b. (Integral a, Num b) => a -> b
fromIntegral a
k forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
log Double
lambda forall a. Num a => a -> a -> a
- Double
k_fac_ln)
  where
    k_fac_ln :: Double
k_fac_ln = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl forall a. Num a => a -> a -> a
(+) Double
0 (forall a b. (a -> b) -> [a] -> [b]
map (forall a. Floating a => a -> a
log forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral) [a
1..a
k])
    lambda :: Double
lambda   = forall a b. (Real a, Fractional b) => a -> b
realToFrac b
mu

fractionalPoisson :: (Num a, Distribution (Poisson b) Integer) => b -> RVarT m a
fractionalPoisson :: forall a b (m :: * -> *).
(Num a, Distribution (Poisson b) Integer) =>
b -> RVarT m a
fractionalPoisson b
mu = forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall a. Num a => Integer -> a
fromInteger (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 :: forall b a.
(CDF (Poisson b) Integer, RealFrac a) =>
b -> a -> Double
fractionalPoissonCDF b
mu a
k = forall (d :: * -> *) t. CDF d t => d t -> t -> Double
cdf (forall b a. b -> Poisson b a
Poisson b
mu) (forall a b. (RealFrac a, Integral b) => a -> b
floor a
k :: Integer)

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

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

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

newtype Poisson b a = Poisson b

instance (RealFloat b, Distribution StdUniform b, Distribution (Erlang Integer) b, Distribution (Binomial b) Integer) => Distribution (Poisson b) Integer where
    rvarT :: forall (n :: * -> *). Poisson b Integer -> RVarT n Integer
rvarT (Poisson b
mu) = forall a b (m :: * -> *).
(Integral a, RealFloat b, Distribution StdUniform b,
 Distribution (Erlang a) b, Distribution (Binomial b) a) =>
b -> RVarT m a
integralPoisson b
mu
instance (Real b, Distribution (Poisson b) Integer) => CDF (Poisson b) Integer where
    cdf :: Poisson b Integer -> Integer -> Double
cdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonCDF b
mu
instance (Real b, Distribution (Poisson b) Integer) => PDF (Poisson b) Integer where
    pdf :: Poisson b Integer -> Integer -> Double
pdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonPDF b
mu
instance (RealFloat b, Distribution StdUniform b, Distribution (Erlang Int) b, Distribution (Binomial b) Int) => Distribution (Poisson b) Int where
    rvarT :: forall (n :: * -> *). Poisson b Int -> RVarT n Int
rvarT (Poisson b
mu) = forall a b (m :: * -> *).
(Integral a, RealFloat b, Distribution StdUniform b,
 Distribution (Erlang a) b, Distribution (Binomial b) a) =>
b -> RVarT m a
integralPoisson b
mu
instance (Real b, Distribution (Poisson b) Int) => CDF (Poisson b) Int where
    cdf :: Poisson b Int -> Int -> Double
cdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonCDF b
mu
instance (Real b, Distribution (Poisson b) Int) => PDF (Poisson b) Int where
    pdf :: Poisson b Int -> Int -> Double
pdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonPDF b
mu
instance (RealFloat b, Distribution StdUniform b, Distribution (Erlang Int8) b, Distribution (Binomial b) Int8) => Distribution (Poisson b) Int8 where
    rvarT :: forall (n :: * -> *). Poisson b Int8 -> RVarT n Int8
rvarT (Poisson b
mu) = forall a b (m :: * -> *).
(Integral a, RealFloat b, Distribution StdUniform b,
 Distribution (Erlang a) b, Distribution (Binomial b) a) =>
b -> RVarT m a
integralPoisson b
mu
instance (Real b, Distribution (Poisson b) Int8) => CDF (Poisson b) Int8 where
    cdf :: Poisson b Int8 -> Int8 -> Double
cdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonCDF b
mu
instance (Real b, Distribution (Poisson b) Int8) => PDF (Poisson b) Int8 where
    pdf :: Poisson b Int8 -> Int8 -> Double
pdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonPDF b
mu
instance (RealFloat b, Distribution StdUniform b, Distribution (Erlang Int16) b, Distribution (Binomial b) Int16) => Distribution (Poisson b) Int16 where
    rvarT :: forall (n :: * -> *). Poisson b Int16 -> RVarT n Int16
rvarT (Poisson b
mu) = forall a b (m :: * -> *).
(Integral a, RealFloat b, Distribution StdUniform b,
 Distribution (Erlang a) b, Distribution (Binomial b) a) =>
b -> RVarT m a
integralPoisson b
mu
instance (Real b, Distribution (Poisson b) Int16) => CDF (Poisson b) Int16 where
    cdf :: Poisson b Int16 -> Int16 -> Double
cdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonCDF b
mu
instance (Real b, Distribution (Poisson b) Int16) => PDF (Poisson b) Int16 where
    pdf :: Poisson b Int16 -> Int16 -> Double
pdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonPDF b
mu
instance (RealFloat b, Distribution StdUniform b, Distribution (Erlang Int32) b, Distribution (Binomial b) Int32) => Distribution (Poisson b) Int32 where
    rvarT :: forall (n :: * -> *). Poisson b Int32 -> RVarT n Int32
rvarT (Poisson b
mu) = forall a b (m :: * -> *).
(Integral a, RealFloat b, Distribution StdUniform b,
 Distribution (Erlang a) b, Distribution (Binomial b) a) =>
b -> RVarT m a
integralPoisson b
mu
instance (Real b, Distribution (Poisson b) Int32) => CDF (Poisson b) Int32 where
    cdf :: Poisson b Int32 -> Int32 -> Double
cdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonCDF b
mu
instance (Real b, Distribution (Poisson b) Int32) => PDF (Poisson b) Int32 where
    pdf :: Poisson b Int32 -> Int32 -> Double
pdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonPDF b
mu
instance (RealFloat b, Distribution StdUniform b, Distribution (Erlang Int64) b, Distribution (Binomial b) Int64) => Distribution (Poisson b) Int64 where
    rvarT :: forall (n :: * -> *). Poisson b Int64 -> RVarT n Int64
rvarT (Poisson b
mu) = forall a b (m :: * -> *).
(Integral a, RealFloat b, Distribution StdUniform b,
 Distribution (Erlang a) b, Distribution (Binomial b) a) =>
b -> RVarT m a
integralPoisson b
mu
instance (Real b, Distribution (Poisson b) Int64) => CDF (Poisson b) Int64 where
    cdf :: Poisson b Int64 -> Int64 -> Double
cdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonCDF b
mu
instance (Real b, Distribution (Poisson b) Int64) => PDF (Poisson b) Int64 where
    pdf :: Poisson b Int64 -> Int64 -> Double
pdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonPDF b
mu
instance (RealFloat b, Distribution StdUniform b, Distribution (Erlang Word) b, Distribution (Binomial b) Word) => Distribution (Poisson b) Word where
    rvarT :: forall (n :: * -> *). Poisson b Word -> RVarT n Word
rvarT (Poisson b
mu) = forall a b (m :: * -> *).
(Integral a, RealFloat b, Distribution StdUniform b,
 Distribution (Erlang a) b, Distribution (Binomial b) a) =>
b -> RVarT m a
integralPoisson b
mu
instance (Real b, Distribution (Poisson b) Word) => CDF (Poisson b) Word where
    cdf :: Poisson b Word -> Word -> Double
cdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonCDF b
mu
instance (Real b, Distribution (Poisson b) Word) => PDF (Poisson b) Word where
    pdf :: Poisson b Word -> Word -> Double
pdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonPDF b
mu
instance (RealFloat b, Distribution StdUniform b, Distribution (Erlang Word8) b, Distribution (Binomial b) Word8) => Distribution (Poisson b) Word8 where
    rvarT :: forall (n :: * -> *). Poisson b Word8 -> RVarT n Word8
rvarT (Poisson b
mu) = forall a b (m :: * -> *).
(Integral a, RealFloat b, Distribution StdUniform b,
 Distribution (Erlang a) b, Distribution (Binomial b) a) =>
b -> RVarT m a
integralPoisson b
mu
instance (Real b, Distribution (Poisson b) Word8) => CDF (Poisson b) Word8 where
    cdf :: Poisson b Word8 -> Word8 -> Double
cdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonCDF b
mu
instance (Real b, Distribution (Poisson b) Word8) => PDF (Poisson b) Word8 where
    pdf :: Poisson b Word8 -> Word8 -> Double
pdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonPDF b
mu
instance (RealFloat b, Distribution StdUniform b, Distribution (Erlang Word16) b, Distribution (Binomial b) Word16) => Distribution (Poisson b) Word16 where
    rvarT :: forall (n :: * -> *). Poisson b Word16 -> RVarT n Word16
rvarT (Poisson b
mu) = forall a b (m :: * -> *).
(Integral a, RealFloat b, Distribution StdUniform b,
 Distribution (Erlang a) b, Distribution (Binomial b) a) =>
b -> RVarT m a
integralPoisson b
mu
instance (Real b, Distribution (Poisson b) Word16) => CDF (Poisson b) Word16 where
    cdf :: Poisson b Word16 -> Word16 -> Double
cdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonCDF b
mu
instance (Real b, Distribution (Poisson b) Word16) => PDF (Poisson b) Word16 where
    pdf :: Poisson b Word16 -> Word16 -> Double
pdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonPDF b
mu
instance (RealFloat b, Distribution StdUniform b, Distribution (Erlang Word32) b, Distribution (Binomial b) Word32) => Distribution (Poisson b) Word32 where
    rvarT :: forall (n :: * -> *). Poisson b Word32 -> RVarT n Word32
rvarT (Poisson b
mu) = forall a b (m :: * -> *).
(Integral a, RealFloat b, Distribution StdUniform b,
 Distribution (Erlang a) b, Distribution (Binomial b) a) =>
b -> RVarT m a
integralPoisson b
mu
instance (Real b, Distribution (Poisson b) Word32) => CDF (Poisson b) Word32 where
    cdf :: Poisson b Word32 -> Word32 -> Double
cdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonCDF b
mu
instance (Real b, Distribution (Poisson b) Word32) => PDF (Poisson b) Word32 where
    pdf :: Poisson b Word32 -> Word32 -> Double
pdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonPDF b
mu
instance (RealFloat b, Distribution StdUniform b, Distribution (Erlang Word64) b, Distribution (Binomial b) Word64) => Distribution (Poisson b) Word64 where
    rvarT :: forall (n :: * -> *). Poisson b Word64 -> RVarT n Word64
rvarT (Poisson b
mu) = forall a b (m :: * -> *).
(Integral a, RealFloat b, Distribution StdUniform b,
 Distribution (Erlang a) b, Distribution (Binomial b) a) =>
b -> RVarT m a
integralPoisson b
mu
instance (Real b, Distribution (Poisson b) Word64) => CDF (Poisson b) Word64 where
    cdf :: Poisson b Word64 -> Word64 -> Double
cdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonCDF b
mu
instance (Real b, Distribution (Poisson b) Word64) => PDF (Poisson b) Word64 where
    pdf :: Poisson b Word64 -> Word64 -> Double
pdf   (Poisson b
mu) = forall a b. (Integral a, Real b) => b -> a -> Double
integralPoissonPDF b
mu

instance Distribution (Poisson b) Integer => Distribution (Poisson b) Float where
    rvarT :: forall (n :: * -> *). Poisson b Float -> RVarT n Float
rvarT (Poisson b
mu) = forall a b (m :: * -> *).
(Num a, Distribution (Poisson b) Integer) =>
b -> RVarT m a
fractionalPoisson b
mu
instance CDF (Poisson b) Integer          => CDF (Poisson b) Float where
    cdf :: Poisson b Float -> Float -> Double
cdf   (Poisson b
mu) = forall b a.
(CDF (Poisson b) Integer, RealFrac a) =>
b -> a -> Double
fractionalPoissonCDF b
mu
instance PDF (Poisson b) Integer          => PDF (Poisson b) Float where
    pdf :: Poisson b Float -> Float -> Double
pdf   (Poisson b
mu) = forall b a.
(PDF (Poisson b) Integer, RealFrac a) =>
b -> a -> Double
fractionalPoissonPDF b
mu
instance Distribution (Poisson b) Integer => Distribution (Poisson b) Double where
    rvarT :: forall (n :: * -> *). Poisson b Double -> RVarT n Double
rvarT (Poisson b
mu) = forall a b (m :: * -> *).
(Num a, Distribution (Poisson b) Integer) =>
b -> RVarT m a
fractionalPoisson b
mu
instance CDF (Poisson b) Integer          => CDF (Poisson b) Double where
    cdf :: Poisson b Double -> Double -> Double
cdf   (Poisson b
mu) = forall b a.
(CDF (Poisson b) Integer, RealFrac a) =>
b -> a -> Double
fractionalPoissonCDF b
mu
instance PDF (Poisson b) Integer          => PDF (Poisson b) Double where
    pdf :: Poisson b Double -> Double -> Double
pdf   (Poisson b
mu) = forall b a.
(PDF (Poisson b) Integer, RealFrac a) =>
b -> a -> Double
fractionalPoissonPDF b
mu