r2rs-nmath 0.1.1

// Translation of nmath's pnbinom
//
// 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, LogProbability64, Positive64, Probability64, Real64};

use crate::{
    distribution::{
        beta::{log_pbeta, pbeta},
        func::bratio,
        pois::{log_ppois, ppois},
    },
    traits::DPQ,
};

/// The distribution function of the negative binomial distribution.
///
/// x = the number of failures before the n-th success
pub fn pnbinom<R: Into<Real64>, PO: Into<Positive64>, PR: Into<Probability64>>(
    x: R,
    size: PO,
    prob: PR,
    lower_tail: bool,
) -> Probability64 {
    pnbinom_inner(x, size, prob, lower_tail, false).into()
}

/// The distribution function of the negative binomial distribution.
///
/// x = the number of failures before the n-th success
pub fn log_pnbinom<R: Into<Real64>, PO: Into<Positive64>, PR: Into<Probability64>>(
    x: R,
    size: PO,
    prob: PR,
    lower_tail: bool,
) -> LogProbability64 {
    pnbinom_inner(x, size, prob, lower_tail, true).into()
}

fn pnbinom_inner<R: Into<Real64>, PO: Into<Positive64>, PR: Into<Probability64>>(
    x: R,
    size: PO,
    prob: PR,
    lower_tail: bool,
    log: bool,
) -> f64 {
    let mut x = x.into().unwrap();
    let size = size.into().unwrap();
    let prob = prob.into().unwrap();

    if !size.is_finite() || !prob.is_finite() {
        return f64::nan();
    }
    if prob <= 0.0 {
        return f64::nan();
    }

    /* limiting case: point mass at zero */
    if size == 0.0 {
        return if x >= 0.0 {
            f64::dt_1(lower_tail, log)
        } else {
            f64::dt_0(lower_tail, log)
        };
    }

    if x < 0.0 {
        return f64::dt_0(lower_tail, log);
    }
    if !x.is_finite() {
        return f64::dt_1(lower_tail, log);
    }
    x = (x + 1e-7).floor();
    if log {
        log_pbeta(prob, size, x + 1.0, lower_tail).unwrap()
    } else {
        pbeta(prob, size, x + 1.0, lower_tail).unwrap()
    }
}

pub fn pnbinom_mu<R: Into<Real64>, PO1: Into<Positive64>, PO2: Into<Positive64>>(
    x: R,
    size: PO1,
    mu: PO2,
    lower_tail: bool,
) -> Probability64 {
    pnbinom_mu_inner(x, size, mu, lower_tail, false).into()
}

pub fn log_pnbinom_mu<R: Into<Real64>, PO1: Into<Positive64>, PO2: Into<Positive64>>(
    x: R,
    size: PO1,
    mu: PO2,
    lower_tail: bool,
) -> LogProbability64 {
    pnbinom_mu_inner(x, size, mu, lower_tail, true).into()
}

pub fn pnbinom_mu_inner<R: Into<Real64>, PO1: Into<Positive64>, PO2: Into<Positive64>>(
    x: R,
    size: PO1,
    mu: PO2,
    lower_tail: bool,
    log: bool,
) -> f64 {
    let mut x = x.into().unwrap();
    let size = size.into().unwrap();
    let mu = mu.into().unwrap();

    if !mu.is_finite() {
        return f64::nan();
    }

    /* limiting case: point mass at zero */
    if size == 0.0 {
        return if x >= 0.0 {
            f64::dt_1(lower_tail, log)
        } else {
            f64::dt_0(lower_tail, log)
        };
    }

    if x < 0.0 {
        return f64::dt_0(lower_tail, log);
    }
    if !x.is_finite() {
        return f64::dt_1(lower_tail, log);
    }
    if !size.is_finite() {
        // limit case: Poisson
        return if log {
            log_ppois(x, mu, lower_tail).unwrap()
        } else {
            ppois(x, mu, lower_tail).unwrap()
        };
    }

    x = (x + 1e-7).floor();
    /* return
     * pbeta(pr, size, x + 1, lower_tail, log);  pr = size/(size + mu), 1-pr = mu/(size+mu)
     *
     *= pbeta_raw(pr, size, x + 1, lower_tail, log)
     *            x.  pin   qin
     *=  bratio (pin,  qin, x., 1-x., &w, &wc, &ierr, log),  and return w or wc ..
     *=  bratio (size, x+1, pr, 1-pr, &w, &wc, &ierr, log) */
    let mut w = f64::nan();
    let mut wc = f64::nan();
    let mut ierr = 0;
    //====
    bratio(
        size,
        x + 1.0,
        size / (size + mu),
        mu / (size + mu),
        &mut w,
        &mut wc,
        &mut ierr,
        log,
    ); /* -> ./toms708.c */
    if ierr != 0 {
        panic!("Error in bratio {}", ierr);
    }
    if lower_tail {
        w
    } else {
        wc
    }
}