// Copyright 2018 Developers of the Rand project.

// Copyright 2016-2017 The Rust Project Developers.

//

// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or

// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license

// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your

// option. This file may not be copied, modified, or distributed

// except according to those terms.

//! The Poisson distribution `Poisson(λ)`.

use crate::{Distribution, Exp1, Normal, StandardNormal, StandardUniform};

use core::fmt;

use num_traits::{Float, FloatConst};

use rand::Rng;

/// The [Poisson distribution](https://en.wikipedia.org/wiki/Poisson_distribution) `Poisson(λ)`.

///

/// The Poisson distribution is a discrete probability distribution with

/// rate parameter `λ` (`lambda`). It models the number of events occurring in a fixed

/// interval of time or space.

///

/// This distribution has density function:

/// `f(k) = λ^k * exp(-λ) / k!` for `k >= 0`.

///

/// # Plot

///

/// The following plot shows the Poisson distribution with various values of `λ`.

/// Note how the expected number of events increases with `λ`.

///

/// ![Poisson distribution](https://raw.githubusercontent.com/rust-random/charts/main/charts/poisson.svg)

///

/// # Example

///

/// ```

/// use rand_distr::{Poisson, Distribution};

///

/// let poi = Poisson::new(2.0).unwrap();

/// let v: f64 = poi.sample(&mut rand::rng());

/// println!("{} is from a Poisson(2) distribution", v);

/// ```

///

/// # Integer vs FP return type

///

/// This implementation uses floating-point (FP) logic internally.

///

/// Due to the parameter limit <code>λ < [Self::MAX_LAMBDA]</code>, it

/// statistically impossible to sample a value larger [`u64::MAX`]. As such, it

/// is reasonable to cast generated samples to `u64` using `as`:

/// `distr.sample(&mut rng) as u64` (and memory safe since Rust 1.45).

/// Similarly, when `λ < 4.2e9` it can be safely assumed that samples are less

/// than `u32::MAX`.

#[derive(Clone, Copy, Debug, PartialEq)]

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]

pub struct Poisson<F>(Method<F>)

where

    F: Float + FloatConst,

    StandardUniform: Distribution<F>;

/// Error type returned from [`Poisson::new`].

#[derive(Clone, Copy, Debug, PartialEq, Eq)]

pub enum Error {

    /// `lambda <= 0`

    ShapeTooSmall,

    /// `lambda = ∞` or `lambda = nan`

    NonFinite,

    /// `lambda` is too large, see [Poisson::MAX_LAMBDA]

    ShapeTooLarge,

}

impl fmt::Display for Error {

    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {

        f.write_str(match self {

            Error::ShapeTooSmall => "lambda is not positive in Poisson distribution",

            Error::NonFinite => "lambda is infinite or nan in Poisson distribution",

            Error::ShapeTooLarge => {

                "lambda is too large in Poisson distribution, see Poisson::MAX_LAMBDA"

            }

        })

    }

}

#[cfg(feature = "std")]

impl std::error::Error for Error {}

#[derive(Clone, Copy, Debug, PartialEq)]

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]

pub(crate) struct KnuthMethod<F> {

    exp_lambda: F,

}

impl<F: Float> KnuthMethod<F> {

    pub(crate) fn new(lambda: F) -> Self {

        KnuthMethod {

            exp_lambda: (-lambda).exp(),

        }

    }

}

#[derive(Clone, Copy, Debug, PartialEq)]

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]

struct RejectionMethod<F> {

    lambda: F,

    s: F,

    d: F,

    l: F,

    c: F,

    c0: F,

    c1: F,

    c2: F,

    c3: F,

    omega: F,

}

impl<F: Float + FloatConst> RejectionMethod<F> {

    pub(crate) fn new(lambda: F) -> Self {

        let b1 = F::from(1.0 / 24.0).unwrap() / lambda;

        let b2 = F::from(0.3).unwrap() * b1 * b1;

        let c3 = F::from(1.0 / 7.0).unwrap() * b1 * b2;

        let c2 = b2 - F::from(15).unwrap() * c3;

        let c1 = b1 - F::from(6).unwrap() * b2 + F::from(45).unwrap() * c3;

        let c0 = F::one() - b1 + F::from(3).unwrap() * b2 - F::from(15).unwrap() * c3;

        RejectionMethod {

            lambda,

            s: lambda.sqrt(),

            d: F::from(6.0).unwrap() * lambda.powi(2),

            l: (lambda - F::from(1.1484).unwrap()).floor(),

            c: F::from(0.1069).unwrap() / lambda,

            c0,

            c1,

            c2,

            c3,

            omega: F::one() / (F::from(2).unwrap() * F::PI()).sqrt() / lambda.sqrt(),

        }

    }

}

#[derive(Clone, Copy, Debug, PartialEq)]

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]

enum Method<F> {

    Knuth(KnuthMethod<F>),

    Rejection(RejectionMethod<F>),

}

impl<F> Poisson<F>

where

    F: Float + FloatConst,

    StandardUniform: Distribution<F>,

{

    /// Construct a new `Poisson` with the given shape parameter

    /// `lambda`.

    ///

    /// The maximum allowed lambda is [MAX_LAMBDA](Self::MAX_LAMBDA).

    pub fn new(lambda: F) -> Result<Poisson<F>, Error> {

        if !lambda.is_finite() {

            return Err(Error::NonFinite);

        }

        if !(lambda > F::zero()) {

            return Err(Error::ShapeTooSmall);

        }

        // Use the Knuth method only for low expected values

        let method = if lambda < F::from(12.0).unwrap() {

            Method::Knuth(KnuthMethod::new(lambda))

        } else {

            if lambda > F::from(Self::MAX_LAMBDA).unwrap() {

                return Err(Error::ShapeTooLarge);

            }

            Method::Rejection(RejectionMethod::new(lambda))

        };

        Ok(Poisson(method))

    }

    /// The maximum supported value of `lambda`

    ///

    /// This value was selected such that

    /// `MAX_LAMBDA + 1e6 * sqrt(MAX_LAMBDA) < 2^64 - 1`,

    /// thus ensuring that the probability of sampling a value larger than

    /// `u64::MAX` is less than 1e-1000.

    ///

    /// Applying this limit also solves

    /// [#1312](https://github.com/rust-random/rand/issues/1312).

    pub const MAX_LAMBDA: f64 = 1.844e19;

}

impl<F> Distribution<F> for KnuthMethod<F>

where

    F: Float + FloatConst,

    StandardUniform: Distribution<F>,

{

    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F {

        let mut result = F::one();

        let mut p = rng.random::<F>();

        while p > self.exp_lambda {

            p = p * rng.random::<F>();

            result = result + F::one();

        }

        result - F::one()

    }

}

impl<F> Distribution<F> for RejectionMethod<F>

where

    F: Float + FloatConst,

    StandardUniform: Distribution<F>,

    StandardNormal: Distribution<F>,

    Exp1: Distribution<F>,

{

    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F {

        // The algorithm is based on:

        // J. H. Ahrens and U. Dieter. 1982.

        // Computer Generation of Poisson Deviates from Modified Normal Distributions.

        // ACM Trans. Math. Softw. 8, 2 (June 1982), 163–179. https://doi.org/10.1145/355993.355997

        // Step F

        let f = |k: F| {

            const FACT: [f64; 10] = [

                1.0, 1.0, 2.0, 6.0, 24.0, 120.0, 720.0, 5040.0, 40320.0, 362880.0,

            ]; // factorial of 0..10

            const A: [f64; 10] = [

                -0.5000000002,

                0.3333333343,

                -0.2499998565,

                0.1999997049,

                -0.1666848753,

                0.1428833286,

                -0.1241963125,

                0.1101687109,

                -0.1142650302,

                0.1055093006,

            ]; // coefficients from Table 1

            let (px, py) = if k < F::from(10.0).unwrap() {

                let px = -self.lambda;

                let py = self.lambda.powf(k) / F::from(FACT[k.to_usize().unwrap()]).unwrap();

                (px, py)

            } else {

                let delta = (F::from(12.0).unwrap() * k).recip();

                let delta = delta - F::from(4.8).unwrap() * delta.powi(3);

                let v = (self.lambda - k) / k;

                let px = if v.abs() <= F::from(0.25).unwrap() {

                    k * v.powi(2)

                        * A.iter()

                            .rev()

                            .fold(F::zero(), |acc, &a| {

                                acc * v + F::from(a).unwrap()

                            }) // Σ a_i * v^i

                        - delta

                } else {

                    k * (F::one() + v).ln() - (self.lambda - k) - delta

                };

                let py = F::one() / (F::from(2.0).unwrap() * F::PI()).sqrt() / k.sqrt();

                (px, py)

            };

            let x = (k - self.lambda + F::from(0.5).unwrap()) / self.s;

            let fx = -F::from(0.5).unwrap() * x * x;

            let fy =

                self.omega * (((self.c3 * x * x + self.c2) * x * x + self.c1) * x * x + self.c0);

            (px, py, fx, fy)

        };

        // Step N

        let normal = Normal::new(self.lambda, self.s).unwrap();

        let g = normal.sample(rng);

        if g >= F::zero() {

            let k1 = g.floor();

            // Step I

            if k1 >= self.l {

                return k1;

            }

            // Step S

            let u: F = rng.random();

            if self.d * u >= (self.lambda - k1).powi(3) {

                return k1;

            }

            let (px, py, fx, fy) = f(k1);

            if fy * (F::one() - u) <= py * (px - fx).exp() {

                return k1;

            }

        }

        loop {

            // Step E

            let e = Exp1.sample(rng);

            let u: F = rng.random() * F::from(2.0).unwrap() - F::one();

            let t = F::from(1.8).unwrap() + e * u.signum();

            if t > F::from(-0.6744).unwrap() {

                let k2 = (self.lambda + self.s * t).floor();

                let (px, py, fx, fy) = f(k2);

                // Step H

                if self.c * u.abs() <= py * (px + e).exp() - fy * (fx + e).exp() {

                    return k2;

                }

            }

        }

    }

}

impl<F> Distribution<F> for Poisson<F>

where

    F: Float + FloatConst,

    StandardUniform: Distribution<F>,

    StandardNormal: Distribution<F>,

    Exp1: Distribution<F>,

{

    #[inline]

    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F {

        match &self.0 {

            Method::Knuth(method) => method.sample(rng),

            Method::Rejection(method) => method.sample(rng),

        }

    }

}

#[cfg(test)]

mod test {

    use super::*;

    #[test]

    #[should_panic]

    fn test_poisson_invalid_lambda_zero() {

        Poisson::new(0.0).unwrap();

    }

    #[test]

    #[should_panic]

    fn test_poisson_invalid_lambda_infinity() {

        Poisson::new(f64::INFINITY).unwrap();

    }

    #[test]

    #[should_panic]

    fn test_poisson_invalid_lambda_neg() {

        Poisson::new(-10.0).unwrap();

    }

    #[test]

    fn poisson_distributions_can_be_compared() {

        assert_eq!(Poisson::new(1.0), Poisson::new(1.0));

    }

}