r2rs-nmath 0.1.1

// Translation of nmath's rnbinom
//
// Copyright (C) 2020 neek-sss
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

use num_traits::Float;
use strafe_type::{FloatConstraint, Positive64, Probability64, Real64};

use crate::{
    distribution::{gamma::rgamma, pois::rpois},
    traits::RNG,
};

/// Random variates from the negative binomial distribution.
///
/// x = the number of failures before the n-th success
///
/// Generate lambda as gamma with shape parameter n and scale
/// parameter $ p/(1-p) $.  Return a Poisson deviate with mean lambda.
pub fn rnbinom<PO: Into<Positive64>, PR: Into<Probability64>, R: RNG>(
    size: PO,
    prob: PR,
    rng: &mut R,
) -> Real64 {
    let mut size = size.into().unwrap();
    let prob = prob.into().unwrap();

    if !prob.is_finite() || size <= 0.0 || prob <= 0.0 {
        /* prob = 1 is ok, PR#1218 */
        return f64::nan().into();
    }
    if !size.is_finite() {
        size = f64::MAX / 2.0
    }
    if prob == 1.0 {
        0.0.into()
    } else {
        let a = size;
        let scale = (1.0 - prob) / prob;
        let mu = rgamma(a, scale, rng).unwrap();
        rpois(mu, rng)
    }
}

pub fn rnbinom_mu<PO1: Into<Positive64>, PO2: Into<Positive64>, R: RNG>(
    size: PO1,
    mu: PO2,
    rng: &mut R,
) -> Real64 {
    let mut size = size.into().unwrap();
    let mu = mu.into().unwrap();

    if !mu.is_finite() || size <= 0.0 {
        return f64::nan().into();
    }
    if !size.is_finite() {
        size = f64::MAX / 2.0
    }

    if mu == 0.0 {
        0.0.into()
    } else {
        let mu = rgamma(size, mu / size, rng).unwrap();
        rpois(mu, rng)
    }
}