r2rs-nmath 0.1.1

// Translation of nmath's pnbeta
//
// 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 nonstdfloat::f128;
use num_traits::{Float, ToPrimitive, Zero};
use strafe_type::{
    FloatConstraint, LogProbability64, Positive64, Probability64, Rational64, Real64,
};

use crate::{
    distribution::func::bratio,
    func::{lbeta, lgamma},
    traits::DPQ,
};

/// Returns the cumulative probability of x for the non-central
/// beta distribution with parameters a, b and non-centrality ncp.
pub fn pnbeta_raw(x: f64, o_x: f64, a: f64, b: f64, ncp: f64) -> f128 {
    /* o_x  == 1 - x  but maybe more accurate */

    /* change errmax and itrmax if desired;
     * original (AS 226, R84) had  (errmax; itrmax) = (1e-6; 100) */
    static errmax: f64 = 1.0e-9; /* 100 is not enough for pf(ncp=200)
                                 see PR#11277 */

    let itrmax = 10000;
    let mut a0 = 0.0;
    let mut lBeta = 0.0;
    let mut c = 0.0;
    let mut errbd = 0.0;
    let mut x0 = 0.0;

    let mut ans = f128::zero();
    let mut ax = f128::zero();
    let mut gx = f128::zero();
    let mut q = f128::zero();
    let mut sumq = f128::zero();

    if ncp < 0.0 || a <= 0.0 || b <= 0.0 {
        return f128::new(f64::nan());
    }

    if x < 0.0 || o_x > 1.0 || (x == 0.0 && o_x == 1.0) {
        return f128::new(0.0);
    }
    if x > 1.0 || o_x < 0.0 || (x == 1.0 && o_x == 0.0) {
        return f128::new(1.0);
    }

    c = ncp / 2.0;

    /* initialize the series */

    x0 = (c - 7.0 * c.sqrt()).max(0.0).floor();
    a0 = a + x0;
    lBeta = lbeta(a0, b).unwrap(); // = tof64!(lgamma(r64!(a0))) + tof64!(lgamma(r64!(b))) - tof64!(lgamma(r64!(a0 + b)));
                                   /* temp = pbeta_raw(x, a0, b, true, false), but using (x, o_x): */

    let mut temp = f64::nan();
    let mut w1 = f64::nan();
    let mut ierr = 0;

    bratio(a0, b, x, o_x, &mut temp, &mut w1, &mut ierr, false); /* -> ./toms708.c */

    if ierr != 0 {
        panic!("Error in bratio {}", ierr);
    }

    gx = f128::new(
        (a0 * x.ln() + b * (if x < 0.5 { (-x).ln_1p() } else { o_x.ln() }) - lBeta - a0.ln()).exp(),
    );
    if a0 > a {
        // x0 >= 1 (and *not* x0 << a)
        q = f128::new((-c + x0 * c.ln() - lgamma(x0 + 1.0).unwrap()).exp())
    } else {
        // a0 = a  <==  x0 << a
        q = f128::new((-c).exp())
    }

    sumq = f128::new(1.0) - q;
    ax = q * f128::new(temp);
    ans = ax;

    /* recurse over subsequent terms until convergence is achieved */
    let mut j = x0.floor(); // x0 could be billions, and is in package EnvStats
    loop {
        j += 1.0;
        temp -= gx.to_f64().unwrap();
        gx *= f128::new(x * (a + b + j - 1.0) / (a + j));
        q *= f128::new(c / j);
        sumq -= q;
        ax = f128::new(temp) * q;
        ans += ax;
        errbd = ((f128::new(temp) - gx) * sumq).to_f64().unwrap();
        if !(errbd > errmax && j < itrmax as f64 + x0) {
            break;
        }
    }

    if errbd > errmax {
        warn!("full precision may not have been achieved in pnbeta");
    }
    if j >= itrmax as f64 + x0 {
        warn!("convergence failed in pnbeta");
    }

    // A small amount of rounding in order to keep out some obviously incorrect values (bug in original)
    if ans.is_sign_negative() && ans.to_f64().unwrap().abs() < 1e30 {
        f128::zero()
    } else {
        ans
    }
}

pub fn pnbeta2(x: f64, o_x: f64, a: f64, b: f64, ncp: f64, lower_tail: bool, log: bool) -> f64 {
    /* o_x  == 1 - x  but maybe more accurate */
    let mut ans = pnbeta_raw(x, o_x, a, b, ncp);

    /* return R_DT_val(ans), but we want to warn about cancellation here */
    if lower_tail {
        if log { ans.ln() } else { ans }.to_f64().unwrap()
    } else {
        if ans > f128::new(1.0 - 1e-10) {
            warn!("full precision may not have been achieved in pnbeta");
        }
        if ans > f128::new(1.0) {
            ans = f128::new(1.0)
        }
        if log {
            (-ans).ln_1p()
        } else {
            (f128::new(1.0)) - ans
        }
        .to_f64()
        .unwrap()
    }
}

pub fn pnbeta<
    RE: Into<Real64>,
    RA1: Into<Rational64>,
    RA2: Into<Rational64>,
    P: Into<Positive64>,
>(
    x: RE,
    a: RA1,
    b: RA2,
    ncp: P,
    lower_tail: bool,
) -> Probability64 {
    pnbeta_inner(x, a, b, ncp, lower_tail, false).into()
}

pub fn log_pnbeta<
    RE: Into<Real64>,
    RA1: Into<Rational64>,
    RA2: Into<Rational64>,
    P: Into<Positive64>,
>(
    x: RE,
    a: RA1,
    b: RA2,
    ncp: P,
    lower_tail: bool,
) -> LogProbability64 {
    pnbeta_inner(x, a, b, ncp, lower_tail, true).into()
}

fn pnbeta_inner<
    RE: Into<Real64>,
    RA1: Into<Rational64>,
    RA2: Into<Rational64>,
    P: Into<Positive64>,
>(
    x: RE,
    a: RA1,
    b: RA2,
    ncp: P,
    lower_tail: bool,
    log: bool,
) -> f64 {
    let x = x.into().unwrap();
    let a = a.into().unwrap();
    let b = b.into().unwrap();
    let ncp = ncp.into().unwrap();

    if let Some(ret) = x.p_bounds_01(0.0, 1.0, lower_tail, log) {
        return ret;
    }
    pnbeta2(x, 1.0 - x, a, b, ncp, lower_tail, log)
}