r2rs-nmath 0.1.1

// Translation of nmath's qgamma
//
// 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, Rational64, Real64,
};

use crate::{
    distribution::{
        gamma::{dgamma, lgamma1p, log_pgamma, pgamma_raw},
        norm::{log_qnorm, qnorm},
    },
    func::lgamma,
    traits::DPQ,
};

pub fn qchisq_appr(p: f64, nu: f64, g: f64, lower_tail: bool, log: bool, tol: f64) -> f64
/* EPS1 */ {
    const C7: f64 = 4.67;
    const C8: f64 = 6.66;
    const C9: f64 = 6.73;
    const C10: f64 = 13.32;

    let mut alpha = 0.0;
    let mut a = 0.0;
    let mut c = 0.0;
    let mut ch = 0.0;
    let mut p1 = 0.0;
    let mut p2 = 0.0;
    let mut q = 0.0;
    let mut t = 0.0;
    let mut x = 0.0;

    /* test arguments and initialise */

    if p.is_nan() || nu.is_nan() {
        return p + nu;
    } /* = [pq]gamma() shape */
    if nu <= 0.0 {
        return f64::nan();
    }

    alpha = 0.5 * nu;
    c = alpha - 1.0;

    p1 = p.dt_log(lower_tail, log);
    if nu < -1.24 * p1 {
        /* for small chi-squared */
        /* log(alpha) + g = log(alpha) + log(gamma(alpha)) =
         *        = log(alpha*gamma(alpha)) = lgamma(alpha+1) suffers from
         *  catastrophic cancellation when alpha << 1
         */
        let lgam1pa = if alpha < 0.5 {
            lgamma1p(alpha)
        } else {
            (alpha.ln()) + g
        }; /* "small nu" : 1.24*(-log(p)) <= nu <= 0.32 */
        ch = ((lgam1pa + p1) / alpha + std::f64::consts::LN_2).exp()
    } else if nu > 0.32 {
        /*  using Wilson and Hilferty estimate */

        x = if log {
            log_qnorm(p, 0.0, 1.0, lower_tail).unwrap()
        } else {
            qnorm(p, 0.0, 1.0, lower_tail).unwrap()
        };
        p1 = 2.0 / (9.0 * nu);
        ch = nu * (x * p1.sqrt() + 1.0 - p1).powf(3.0);

        /* approximation for p tending to 1: */
        if ch > 2.2 * nu + 6.0 {
            ch = -2.0 * (p.dt_clog(lower_tail, log) - c * (0.5 * ch).ln() + g)
        }
    } else {
        /* "small nu" : 1.24*(-log(p)) <= nu <= 0.32 */
        ch = 0.4;
        a = p.dt_clog(lower_tail, log) + g + c * std::f64::consts::LN_2;
        loop {
            q = ch;
            p1 = 1.0 / (1.0 + ch * (C7 + ch));
            p2 = ch * (C9 + ch * (C8 + ch));
            t = -0.5 + (C7 + 2.0 * ch) * p1 - (C9 + ch * (C10 + 3.0 * ch)) / p2;
            ch -= (1.0 - (a + 0.5 * ch).exp() * p2 * p1) / t;
            if !((q - ch).abs() > tol * ch.abs()) {
                break;
            }
        }
    }
    return ch;
}

/// Compute the quantile function of the gamma distribution.
///
/// This function is based on the Applied Statistics
/// Algorithm AS 91 ("ppchi2") and via pgamma(.) AS 239.
///
/// R core improvements:
/// o  lower_tail, log
/// o  non-trivial result for p outside [0.000002, 0.999998]
/// o  p ~ 1 no longer gives +Inf; final Newton step(s)
pub fn qgamma<PR: Into<Probability64>, PO: Into<Positive64>, R: Into<Rational64>>(
    p: PR,
    alpha: PO,
    scale: R,
    lower_tail: bool,
) -> Real64 {
    let p = p.into().unwrap();
    qgamma_inner(p, alpha, scale, lower_tail, false)
}

/// Compute the quantile function of the gamma distribution.
///
/// This function is based on the Applied Statistics
/// Algorithm AS 91 ("ppchi2") and via pgamma(.) AS 239.
///
/// R core improvements:
/// o  lower_tail, log
/// o  non-trivial result for p outside [0.000002, 0.999998]
/// o  p ~ 1 no longer gives +Inf; final Newton step(s)
pub fn log_qgamma<LP: Into<LogProbability64>, P: Into<Positive64>, R: Into<Rational64>>(
    p: LP,
    alpha: P,
    scale: R,
    lower_tail: bool,
) -> Real64 {
    let p = p.into().unwrap();
    qgamma_inner(p, alpha, scale, lower_tail, true)
}

pub fn qgamma_inner<P: Into<Positive64>, R: Into<Rational64>>(
    mut p: f64,
    alpha: P,
    scale: R,
    lower_tail: bool,
    log: bool,
) -> Real64 {
    /*   shape = alpha */
    let alpha = alpha.into().unwrap();
    let scale = scale.into().unwrap();

    const EPS1: f64 = 1e-2;
    const EPS2: f64 = 5e-7; /* final precision of AS 91 */
    const EPS_N: f64 = 1e-15; /* precision of Newton step / iterations */
    // const LN_EPS: f64 = -36.043653389117156; /* = log(.Machine$double.eps) iff IEEE_754 */
    const MAXIT: i32 = 1000; /* was 20 */

    const P_MIN: f64 = 1e-100; /* was 0.000002 = 2e-6 */
    const P_MAX: f64 = 1.0 - 1e-14; /* was (1-1e-12) and 0.999998 = 1 - 2e-6 */

    static i420: f64 = 1.0 / 420.0;
    static i2520: f64 = 1.0 / 2520.0;
    static i5040: f64 = 1.0 / 5040.0;

    let mut p_ = 0.0;
    let mut a = 0.0;
    let mut b = 0.0;
    let mut c = 0.0;
    let mut g = 0.0;
    let mut ch = 0.0;
    let mut ch0 = 0.0;
    let mut p1 = 0.0;
    let mut p2 = 0.0;
    let mut q = 0.0;
    let mut s1 = 0.0;
    let mut s2 = 0.0;
    let mut s3 = 0.0;
    let mut s4 = 0.0;
    let mut s5 = 0.0;
    let mut s6 = 0.0;
    let mut t = 0.0;
    let mut x = 0.0;
    let mut i = 0;
    let mut max_it_Newton = 1;

    /* test arguments and initialise */

    if let Some(ret) = p.q_p01_boundaries(0.0, f64::infinity(), lower_tail, log) {
        return ret.into();
    }

    if alpha == 0.0 {
        /* all mass at 0 : */
        return 0.0.into();
    }

    if alpha < 1e-10 {
        /* Warning seems unnecessary now: */
        max_it_Newton = 7
        /* may still be increased below */
    } /* lower_tail prob (in any case) */

    p_ = p.dt_qiv(lower_tail, log);

    g = lgamma(alpha).unwrap(); /* log Gamma(v/2) */

    /*----- Phase I : Starting Approximation */
    ch = qchisq_appr(p, 2.0 * alpha, g, lower_tail, log, EPS1);
    if !ch.is_finite() {
        /* forget about all iterations! */
        max_it_Newton = 0
    } else if ch < 5e-7 || (log && (p_ > P_MAX || p_ < P_MIN)) {
        /* Corrected according to AS 91; MM, May 25, 1999 */
        max_it_Newton = 20
        /* and do Newton steps */
    } else {
        /*----- Phase II: Iteration
         * Call pgamma() [AS 239] and calculate seven term taylor series
         */
        c = alpha - 1.0; /* used below, is "const" */
        s6 = (120.0 + c * (346.0 + 127.0 * c)) * i5040; /* save initial approx. */
        ch0 = ch; /*was  return ML_NAN;*/
        i = 1;
        while i <= MAXIT {
            q = ch;
            p1 = 0.5 * ch;
            // TODO: Handling of this error (should it be done like this?)
            if log && p1.is_nan() {
                p2 = f64::nan();
            } else {
                p2 = p_ - pgamma_raw(p1, alpha, true, false);
            }
            if !p2.is_finite() || ch <= 0.0 {
                ch = ch0;
                max_it_Newton = 27;
                break;
            } else {
                t = p2 * (alpha * std::f64::consts::LN_2 + g + p1 - c * ch.ln()).exp();
                b = t / ch;
                a = 0.5 * t - b * c;
                s1 =
                    (210.0 + a * (140.0 + a * (105.0 + a * (84.0 + a * (70.0 + 60.0 * a))))) * i420;
                s2 = (420.0 + a * (735.0 + a * (966.0 + a * (1141.0 + 1278.0 * a)))) * i2520;
                s3 = (210.0 + a * (462.0 + a * (707.0 + 932.0 * a))) * i2520;
                s4 = (252.0 + a * (672.0 + 1182.0 * a) + c * (294.0 + a * (889.0 + 1740.0 * a)))
                    * i5040;
                s5 = (84.0 + 2264.0 * a + c * (1175.0 + 606.0 * a)) * i2520;
                ch += t
                    * (1.0 + 0.5 * t * s1
                        - b * c * (s1 - b * (s2 - b * (s3 - b * (s4 - b * (s5 - b * s6))))));
                if (q - ch).abs() < EPS2 * ch {
                    break;
                }
                if (q - ch).abs() > 0.1 * ch {
                    /* diverging? -- also forces ch > 0 */
                    if ch < q {
                        ch = 0.9 * q
                    } else {
                        ch = 1.1 * q
                    }
                }
                i += 1
            }
        }
    }
    /* no convergence in MAXIT iterations -- but we add Newton now... */
    /* was
     *    ML_ERROR(ME_PRECISION, "qgamma");
     * does nothing in R !*/

    /* PR# 2214 :  From: Morten Welinder <terra@diku.dk>, Fri, 25 Oct 2002 16:50
    --------  To: R-bugs@biostat.ku.dk     Subject: qgamma precision

    * With a final Newton step, double accuracy, e.g. for (p= 7e-4; nu= 0.9)
    *
    * Improved (MM): - only if rel.Err > EPS_N (= 1e-15);
    *      - also for lower_tail = false  or log = true
    *      - optionally *iterate* Newton
    */
    x = 0.5 * scale * ch;
    if max_it_Newton != 0 {
        /* always use log scale */
        if !log {
            p = p.ln();
        }
        if x == 0.0 {
            let _1_p = 1.0 + 1e-7;
            let _1_m = 1.0 - 1e-7;
            x = 2.2250738585072014e-308;
            p_ = log_pgamma(x, alpha, scale, lower_tail).unwrap();
            if lower_tail && p_ > p * _1_p || !lower_tail && p_ < p * _1_m {
                return 0.0.into();
            }
            /* else:  continue, using x = DBL_MIN instead of  0  */
        } else {
            p_ = log_pgamma(x, alpha, scale, lower_tail).unwrap()
        } /* PR#14710 */
        if p_ == f64::NEG_INFINITY {
            return 0.0.into();
        }
        i = 1;
        while i <= max_it_Newton {
            p1 = p_ - p;
            if p1.abs() < (EPS_N * p).abs() {
                break;
            }
            /* else */
            g = dgamma(x, alpha, scale, true).unwrap();
            if g == (f64::d_0(true)) {
                break;
            }
            /* else :
             * delta x = f(x)/f'(x);
             * if(log) f(x) := log P(x) - p; f'(x) = d/dx log P(x) = P' / P
             * ==> f(x)/f'(x) = f*P / P' = f*exp(p_) / P' (since p_ = log P(x))
             */
            t = (p1) * (p_ - g).exp(); /* = "delta x" */
            t = if lower_tail { (x) - t } else { (x) + t }; /* else : */
            p_ = log_pgamma(t, alpha, scale, lower_tail).unwrap();
            if (p_ - p).abs() > p1.abs() || i > 1 && (p_ - p).abs() == p1.abs() {
                /* <- against flip-flop */
                break;
            } else {
                x = t;
                i += 1
            }
        }
    }

    x.into()
}