r2rs-nmath 0.1.1

// Translation of nmath's pnchisq
//
// 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 std::f64::consts::LN_10;

use nonstdfloat::f128;
use num_traits::{Float, ToPrimitive, Zero};
use strafe_type::{
    FloatConstraint, LogProbability64, Positive64, Probability64, Rational64, Real64,
};

use crate::{
    distribution::{
        chisq::{log_pchisq, pchisq},
        func::logspace_add,
    },
    func::lgamma,
    traits::DPQ,
};

static _dbl_min_exp: f64 = std::f64::consts::LN_2 * f64::MIN_EXP as f64;

pub fn pnchisq<RE: Into<Real64>, RA: Into<Rational64>, P: Into<Positive64>>(
    x: RE,
    df: RA,
    ncp: P,
    lower_tail: bool,
) -> Probability64 {
    pnchisq_inner(x, df, ncp, lower_tail, false).into()
}

pub fn log_pnchisq<RE: Into<Real64>, RA: Into<Rational64>, P: Into<Positive64>>(
    x: RE,
    df: RA,
    ncp: P,
    lower_tail: bool,
) -> LogProbability64 {
    pnchisq_inner(x, df, ncp, lower_tail, true).into()
}

fn pnchisq_inner<RE: Into<Real64>, RA: Into<Rational64>, P: Into<Positive64>>(
    x: RE,
    df: RA,
    ncp: P,
    lower_tail: bool,
    log: bool,
) -> f64 {
    let mut x = x.into().unwrap();
    let df = df.into().unwrap();
    let ncp = ncp.into().unwrap();

    if x == 0.0 {
        x = f64::min_positive_value();
    }

    let mut ans = 0.0;
    if !df.is_finite() || !ncp.is_finite() {
        return f64::nan();
    }

    ans = pnchisq_raw(
        x,
        df,
        ncp,
        1e-12,
        8.0 * 2.2204460492503131e-16,
        1000000,
        lower_tail,
        log,
    );

    if x <= 0. || x == f64::infinity() {
        return ans; // because it's perfect
    }

    if ncp >= 80.0 {
        if lower_tail {
            ans = ans.min(f64::d_1(log))
            /* e.g., pchisq(555, 1.01, ncp = 80) */
        } else {
            /* since we computed the other tail cancellation is likely */
            // FIXME: There are cases where  ans == 0. if(!log_p) is perfect
            if ans < if log { -10. * LN_10 } else { 1e-10 } {
                warn!("full precision may not have been achieved in pnchisq");
            }
            if !log && ans < 0.0 {
                ans = 0.0;
            }
            /* Precaution PR#7099 */
        }
    }

    /* MM: the following "hack" from c51179 (<--> PR#14216, by Jerry Lewis)
     * -- is "kind of ok" ... but potentially suboptimal: we do  log1p(- p(*, <other tail>, log=FALSE)),
     *    but that  p(*, log=FALSE) may already be an exp(.) or even expm1(..)
     *   <---> "in principle"  this check should happen there, not here  */

    if !log || ans < -1e-8 {
        ans
    } else {
        // log_p (==> ans <= 0) &&  -1e-8 <= ans <= 0
        // prob. = ans.exp() is near one: we can do better using the other tail
        ans = pnchisq_raw(
            x,
            df,
            ncp,
            1e-12,
            8.0 * 2.2204460492503131e-16,
            1000000,
            !lower_tail,
            false,
        );
        (-ans).ln_1p()
    }
}

pub fn pnchisq_raw(
    x: f64,
    f: f64,
    theta: f64,
    errmax: f64,
    reltol: f64,
    itrmax: i32,
    lower_tail: bool,
    log: bool,
) -> f64 {
    let mut lam = 0.0; /* initialized for -Wall */
    let mut x2 = 0.0;
    let mut f2 = 0.0;
    let mut term = 0.0;
    let mut bound = 0.0;
    let mut f_x_2n = 0.0;
    let mut f_2n = 0.0;
    let mut l_lam = -1.0;
    let mut l_x = -1.0;
    let mut n = 0;
    let mut lamSml = false;
    let mut tSml = false;
    let mut is_r = false;
    let mut is_b = false;
    let mut is_it = false;
    let mut ans = f128::zero();
    let mut u = f128::zero();
    let mut v = f128::zero();
    let mut t = f128::zero();
    let mut lt = f128::zero();
    let mut lu = f128::new(-(1.0));

    if x <= 0.0 {
        if x == 0.0 && f == 0.0 {
            // chi^2_0(.) has point mass at zero
            let _L = (-0.5) * theta;
            // = -lambda
            return if lower_tail {
                _L.d_exp(log)
            } else if log {
                _L.log1_exp()
            } else {
                -_L.exp_m1()
            };
        }
        /* x < 0  or {x==0, f > 0} */
        return f64::dt_0(lower_tail, log);
    }
    if !x.is_finite() {
        return f64::dt_1(lower_tail, log);
    }

    /* This is principally for use from qnchisq */
    // R_CheckUserInterrupt(); // if(theta < 80)
    if theta < 80.0 {
        /* use 110 for Inf, as ppois(110, 80/2, lower.tail=false) is 2e-20 */
        let mut ans_0 = f128::zero();
        let mut i = 0;
        // Have  pgamma(x,s) < x^s / Gamma(s+1) (< and ~= for small x)
        // ==> pchisq(x, f) = pgamma(x, f/2, 2) = pgamma(x/2, f/2)
        // <  (x/2)^(f/2) / Gamma(f/2+1) < eps
        // <==>  f/2 * (x/2).ln() - (Gamma(f/2+1)).ln() < eps.ln() ( ~= -708.3964 )
        // <==>        (x/2).ln() < 2/f*((Gamma(f/2+1)).ln() + eps.ln())
        // <==> x.ln() < 2.ln() + 2/f*((Gamma(f/2+1)).ln() + eps.ln())
        return if lower_tail
            && f > 0.0
            && x.ln()
                < std::f64::consts::LN_2 + 2.0 / f * (lgamma(f / 2.0 + 1.0).unwrap() + _dbl_min_exp)
        {
            // all  pchisq(x, f+2*i, lower_tail, false), i=0,...,110 would underflow to 0.
            // ==> work in log scale
            let lambda = 0.5 * theta; // < 40
            let mut sum = 0.0;
            let mut sum2 = 0.0;
            let mut pr = -lambda;
            let log_lam = lambda.ln();
            sum2 = f64::NEG_INFINITY;
            sum = sum2;
            /* we need to renormalize here: the result could be very close to 1 */
            i = 0; // does this need a feature test?
            while i < 110 {
                sum2 = logspace_add(sum2, pr);
                sum = logspace_add(
                    sum,
                    pr + log_pchisq(x, f + (2 * i) as f64, lower_tail).unwrap(),
                );
                if sum2 >= -1e-15 {
                    break;
                }
                i += 1;
                pr += log_lam - (i as f64).ln()
            }
            ans_0 = f128::new(sum - sum2);
            if log { ans_0 } else { ans_0.exp() }.to_f64().unwrap()
        } else {
            let lambda_0 = f128::new(0.5 * theta); // < 40
            let mut sum_0 = f128::new(0.0);
            let mut sum2_0 = f128::new(0.0);
            let mut pr_0 = (-lambda_0).exp();
            /* we need to renormalize here: the result could be very close to 1 */
            i = 0;
            while i < 110 {
                // pr == (-lambda).exp() lambda^i / i!  ==  dpois(i, lambda)
                sum2_0 += pr_0;
                // pchisq(*, i, *) is  strictly decreasing to 0 for lower_tail=true
                // and strictly increasing to 1 for lower_tail=false
                sum_0 += pr_0 * f128::new(pchisq(x, f + (2 * i) as f64, lower_tail));
                if sum2_0 >= f128::new(1.0 - 1e-15) {
                    break;
                }
                i += 1;
                pr_0 *= lambda_0 / f128::new(i as f64)
            }
            ans_0 = sum_0 / sum2_0;
            if log { ans_0.ln() } else { ans_0 }.to_f64().unwrap()
        };
    }
    // else: theta == ncp >= 80 --------------------------------------------
    // Series expansion ------- FIXME: log=true, lower_tail=false only applied at end

    lam = 0.5 * theta; // = lambda = ncp/2
    lamSml = -lam < _dbl_min_exp;
    if lamSml {
        // originally error: "non centrality parameter too large for current algorithm"
        u = f128::new(0.0); /* == ln(u) */
        lu = f128::new(-lam);
        l_lam = lam.ln()
    } else {
        u = f128::new((-lam).exp())
    }

    /* evaluate the first term */
    v = u;
    x2 = 0.5 * x;
    f2 = 0.5 * f;
    f_x_2n = f - x;

    if f2 * 2.2204460492503131e-16 > 0.125 && {
        t = f128::new(x2 - f2);
        (t.abs()) < f128::new((2.2204460492503131e-16).sqrt() * f2)
    } {
        /* evade cancellation error */
        /* t = ((1 - t)*(2 - t/(f2 + 1))).exp() / (2*M_PI*(f2 + 1)).sqrt();*/
        lt = (f128::new(1.0) - t) * (f128::new(2.0) - t / f128::new(f2 + 1.0))
            - f128::new(strafe_consts::LN_SQRT_2TPI)
            - f128::new(0.5 * (f2 + 1.0).ln())
    } else {
        /* Usual case 2: careful not to overflow .. : */
        lt = f128::new(f2 * x2.ln() - x2 - lgamma(f2 + 1.0).unwrap())
    } /* else */

    tSml = lt < f128::new(_dbl_min_exp);
    if tSml {
        if x > f + theta + 5.0 * (2.0 * (f + 2.0 * theta)).sqrt() {
            /* x > E[X] + 5* sigma(X) */
            return f64::dt_1(lower_tail, log);
            /* FIXME: could be more accurate than 0.0 */
        }
        l_x = x.ln();
        term = 0.0;
        ans = f128::new(term);
        t = f128::new(0.0)
    } else {
        t = lt.exp();
        term = (v * t).to_f64().unwrap();
        ans = f128::new(term)
    }

    n = 1;
    f_2n = f + 2.0;
    f_x_2n += 2.0;
    loop {
        if f_x_2n > 0.0 {
            /* find the error bound and check for convergence */

            bound = (t * f128::new(x) / f128::new(f_x_2n)).to_f64().unwrap();
            is_it = false;
            is_r = is_it;
            /* convergence only if BOTH absolute and relative error < 'bnd' */
            is_b = bound <= errmax;
            if is_b && {
                is_r = f128::new(term) <= f128::new(reltol) * ans;
                is_r
            } || {
                is_it = n > itrmax;
                is_it
            } {
                break;
            }
        }
        /* evaluate the next term of the */
        /* expansion and then the partial sum */

        if lamSml {
            lu += f128::new(l_lam - (n as f64).ln()); /* u = u* lam / n */
            if lu >= f128::new(_dbl_min_exp) {
                /* no underflow anymore ==> change regime */
                u = lu.exp(); /* the first non-0 'u' */
                v = u; /* t <- t * (x / f2n) */
                lamSml = false
            }
        } else {
            u *= f128::new(lam / n as f64);
            v += u
        }
        if tSml {
            lt += f128::new(l_x - f_2n.ln());
            if lt >= f128::new(_dbl_min_exp) {
                /* no underflow anymore ==> change regime */
                t = lt.exp(); /* the first non-0 't' */
                tSml = false
            }
        } else {
            t *= f128::new(x / f_2n)
        }
        if lamSml as u64 == 0 && tSml as u64 == 0 {
            term = (v * t).to_f64().unwrap();
            ans += f128::new(term)
        }
        n += 1;
        f_2n += 2.0;
        f_x_2n += 2.0
    }

    if is_it {
        warn!("pnchisq(x={}, ..): not converged in {} iter.", x, itrmax);
    }
    let dans = ans.to_f64().unwrap();
    return dans.dt_val(lower_tail, log);
}