r2rs-nmath 0.1.1

// Translation of nmath's pgamma
//
// 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::{
        norm::{dnorm, log_pnorm, pnorm},
        pois::dpois_raw,
    },
    func::lgamma,
    traits::DPQ,
};

/// Scalefactor:= (2^32)^8 = 2^256 = 1.157921e+77 */
#[macro_export]
macro_rules! SQR {
    ($x:expr) => {{
        $x * $x
    }};
}
static scalefactor: f64 = SQR!(SQR!(SQR!(4294967296.0)));

/// If |x| > |k| * M_cutoff,  then  log[ (-x).exp() * k^x ]  =~=  -x */
static M_cutoff: f64 = std::f64::consts::LN_2 * f64::MAX_EXP as f64 / 2.2204460492503131e-16;

/// Continued fraction for calculation of
/// $ 1/i + x/(i+d) + x^2/(i+2*d) + x^3/(i+3*d) + ... = sum_{k=0}^\infty x^k/(i+k*d) $
///
/// auxiliary in log1pmx() and lgamma1p()
fn logcf(x: f64, i: f64, d: f64, eps: f64) -> f64
/* ~ relative tolerance */ {
    let mut c1 = 2.0 * d;
    let mut c2 = i + d;
    let mut c4 = c2 + d;
    let mut a1 = c2;
    let mut b1 = i * (c2 - i * x);
    let mut b2 = d * d * x;
    let mut a2 = c4 * c2 - b2;

    b2 = c4 * b1 - i * b2;

    while (a2 * b1 - a1 * b2).abs() > (eps * b1 * b2).abs() {
        let mut c3 = c2 * c2 * x;
        c2 += d;
        c4 += d;
        a1 = c4 * a2 - c3 * a1;
        b1 = c4 * b2 - c3 * b1;

        c3 = c1 * c1 * x;
        c1 += d;
        c4 += d;
        a2 = c4 * a1 - c3 * a2;
        b2 = c4 * b1 - c3 * b2;

        if b2.abs() > scalefactor {
            a1 /= scalefactor;
            b1 /= scalefactor;
            a2 /= scalefactor;
            b2 /= scalefactor
        } else if b2.abs() < 1.0 / scalefactor {
            a1 *= scalefactor;
            b1 *= scalefactor;
            a2 *= scalefactor;
            b2 *= scalefactor
        }
    }

    a2 / b2
}

/// Accurate calculation of log(1+x)-x, particularly for small x.
pub fn log1pmx(x: f64) -> f64 {
    static minLog1Value: f64 = -0.79149064;

    if x > 1.0 || x < minLog1Value {
        x.ln_1p() - x
    } else {
        /* -.791 <=  x <= 1  -- expand in  [x/(2+x)]^2 =: y :
         * (1+x).ln() - x =  x/(2+x) * [ 2 * y * S(y) - x],  with
         * ---------------------------------------------
         * S(y) = 1/3 + y/5 + y^2/7 + ... = \sum_{k=0}^\infty  y^k / (2k + 3)
         */
        let r = x / (2.0 + x);
        let y = r * r;
        if x.abs() < 1e-2 {
            static two: f64 = 2.0;
            r * ((((two / 9.0 * y + two / 7.0) * y + two / 5.0) * y + two / 3.0) * y - x)
        } else {
            static tol_logcf: f64 = 1e-14;
            r * (2.0 * y * logcf(y, 3.0, 2.0, tol_logcf) - x)
        }
    }
}

/// Compute log(gamma(a+1)) accurately also for small a (0 < a < 0.5).
pub fn lgamma1p(a: f64) -> f64 {
    if a.abs() >= 0.5 {
        return lgamma(a + 1.0).unwrap();
    }

    let eulers_const = 0.5772156649015328606065120900824024;

    /* coeffs[i] holds (zeta(i+2)-1)/(i+2) , i = 0:(N-1), N = 40 : */
    let N = 40; /* zeta(N+2)-1 */
    static coeffs: [f64; 40] = [
        0.3224670334241132182362075833230126e-0, /* = (zeta(2)-1)/2 */
        0.6735230105319809513324605383715000e-1, /* = (zeta(3)-1)/3 */
        0.2058080842778454787900092413529198e-1,
        0.7385551028673985266273097291406834e-2,
        0.2890510330741523285752988298486755e-2,
        0.1192753911703260977113935692828109e-2,
        0.5096695247430424223356548135815582e-3,
        0.2231547584535793797614188036013401e-3,
        0.9945751278180853371459589003190170e-4,
        0.4492623673813314170020750240635786e-4,
        0.2050721277567069155316650397830591e-4,
        0.9439488275268395903987425104415055e-5,
        0.4374866789907487804181793223952411e-5,
        0.2039215753801366236781900709670839e-5,
        0.9551412130407419832857179772951265e-6,
        0.4492469198764566043294290331193655e-6,
        0.2120718480555466586923135901077628e-6,
        0.1004322482396809960872083050053344e-6,
        0.4769810169363980565760193417246730e-7,
        0.2271109460894316491031998116062124e-7,
        0.1083865921489695409107491757968159e-7,
        0.5183475041970046655121248647057669e-8,
        0.2483674543802478317185008663991718e-8,
        0.1192140140586091207442548202774640e-8,
        0.5731367241678862013330194857961011e-9,
        0.2759522885124233145178149692816341e-9,
        0.1330476437424448948149715720858008e-9,
        0.6422964563838100022082448087644648e-10,
        0.3104424774732227276239215783404066e-10,
        0.1502138408075414217093301048780668e-10,
        0.7275974480239079662504549924814047e-11,
        0.3527742476575915083615072228655483e-11,
        0.1711991790559617908601084114443031e-11,
        0.8315385841420284819798357793954418e-12,
        0.4042200525289440065536008957032895e-12,
        0.1966475631096616490411045679010286e-12,
        0.9573630387838555763782200936508615e-13,
        0.4664076026428374224576492565974577e-13,
        0.2273736960065972320633279596737272e-13,
        0.1109139947083452201658320007192334e-13, /* = (zeta(40+1)-1)/(40+1) */
    ];

    let c = 0.2273736845824652515226821577978691e-12;
    let tol_logcf = 1e-14;
    let mut lgam = 0.0;
    let mut i: i32 = 0;

    if a.abs() >= 0.5 {
        return lgamma(a + 1.0).unwrap();
    }

    /* Abramowitz & Stegun 6.1.33 : for |x| < 2,
     * <==> (gamma(1+x)).ln() = -((1+x).ln() - x) - gamma*x + x^2 * \sum_{n=0}^\infty c_n (-x)^n
     * where c_n := (Zeta(n+2) - 1)/(n+2)  = coeffs[n]
     *
     * Here, another convergence acceleration trick is used to compute
     * lgam(x) :=  sum_{n=0..Inf} c_n (-x)^n
     */
    lgam = c * logcf(-a / 2.0, (N + 2) as f64, 1.0, tol_logcf);
    i = N - 1;
    while i >= 0 {
        lgam = coeffs[i as usize] - a * lgam;
        i -= 1;
    }

    (a * lgam - eulers_const) * a - log1pmx(a)
}

/// dpois_wrap (x__1, lambda) := dpois(x__1 - 1, lambda);  where
/// dpois(k, L) := (-L).exp() L^k / gamma(k+1)  {the usual Poisson probabilities}
/// and  dpois*(.., give_log = true) :=  ( dpois*(..) ).ln()
fn dpois_wrap(x_plus_1: f64, lambda: f64, log: bool) -> f64 {
    if !lambda.is_finite() {
        return f64::d_0(log);
    }
    if x_plus_1 > 1.0 {
        return dpois_raw(x_plus_1 - 1.0, lambda, log);
    }
    if lambda > (x_plus_1 - 1.0).abs() * M_cutoff {
        (-lambda - lgamma(x_plus_1).unwrap()).d_exp(log)
    } else {
        let d = dpois_raw(x_plus_1, lambda, log);
        if log {
            d + (x_plus_1 / lambda).ln()
        } else {
            d * (x_plus_1 / lambda)
        }
    }
}

/// Abramowitz and Stegun 6.5.29 \[right\]
fn pgamma_smallx(x: f64, alph: f64, lower_tail: bool, log: bool) -> f64 {
    let mut sum = 0.0;
    let mut c = alph;
    let mut n = 0.0;
    let mut term = 0.0;

    /*
     * Relative to 6.5.29 all terms have been multiplied by alph
     * and the first, thus being 1, is omitted.
     */
    loop {
        n += 1.0;
        c *= -x / n;
        term = c / (alph + n);
        sum += term;
        if !(term.abs() > 2.2204460492503131e-16 * sum.abs()) {
            break;
        }
    }

    if lower_tail {
        let f1 = if log { sum.ln_1p() } else { (1.0) + sum };
        let mut f2 = 0.0;
        if alph > 1.0 {
            f2 = dpois_raw(alph, x, log);
            f2 = if log { (f2) + x } else { (f2) * x.exp() }
        } else if log {
            f2 = alph * x.ln() - lgamma1p(alph)
        } else {
            f2 = x.powf(alph) / lgamma1p(alph).exp()
        }
        if log {
            (f1) + f2
        } else {
            (f1) * f2
        }
    } else {
        let lf2 = alph * x.ln() - lgamma1p(alph);
        if log {
            (sum.ln_1p() + lf2).log1_exp()
        } else {
            let f1m1 = sum;
            let f2m1 = lf2.exp_m1();
            -(f1m1 + f2m1 + f1m1 * f2m1)
        }
    }
}

fn pd_upper_series(x: f64, mut y: f64, log: bool) -> f64 {
    let mut term = x / y;
    let mut sum = term;

    loop {
        y += 1.0;
        term *= x / y;
        sum += term;
        if !(term > sum * 2.2204460492503131e-16) {
            break;
        }
    }

    /* sum =  \sum_{n=1}^ oo  x^n     / (y*(y+1)*...*(y+n-1))
     *    =  \sum_{n=0}^ oo  x^(n+1) / (y*(y+1)*...*(y+n))
     *    =  x/y * (1 + \sum_{n=1}^oo x^n / ((y+1)*...*(y+n)))
     *    ~  x/y +  o(x/y)   {which happens when alph -> Inf}
     */
    if log {
        sum.ln()
    } else {
        sum
    }
}

/// Continued fraction for calculation of
/// scaled upper-tail F_{gamma}
/// ~=  (y / d) * [1 +  (1-y)/d +  O( ((1-y)/d)^2 ) ]
fn pd_lower_cf(y: f64, d: f64) -> f64 {
    const maxit: usize = 200000;

    let mut f = 0.0;
    let mut of = 0.0;
    let mut f0 = 0.0;
    let mut i = 0;
    let mut c2 = 0.0;
    let mut c3 = 0.0;
    let mut c4 = 0.0;
    let mut a1 = 0.0;
    let mut b1 = 0.0;
    let mut a2 = 0.0;
    let mut b2 = 0.0;

    if y == 0.0 {
        return 0.0;
    }

    f0 = y / d;
    /* Needed, e.g. for  pgamma(10^c(100,295), shape= 1.1, log=true): */
    if (y - 1.0).abs() < d.abs() * 2.2204460492503131e-16 {
        /* includes y < d = Inf */
        return f0;
    }

    if f0 > 1.0 {
        f0 = 1.0
    }
    c2 = y;
    c4 = d; /* original (y,d), *not* potentially scaled ones!*/

    a1 = 0.0;
    b1 = 1.0;
    a2 = y;
    b2 = d;

    while b2 > scalefactor {
        a1 /= scalefactor;
        b1 /= scalefactor;
        a2 /= scalefactor;
        b2 /= scalefactor
    }

    i = 0;
    of = -1.0; /* far away */
    while i < maxit {
        i += 1;
        c2 -= 1.0;
        c3 = i as f64 * c2;
        c4 += 2.0;
        /* c2 = y - i,  c3 = i(y - i),  c4 = d + 2i,  for i odd */
        a1 = c4 * a2 + c3 * a1;
        b1 = c4 * b2 + c3 * b1;

        i += 1;
        c2 -= 1.0;
        c3 = i as f64 * c2;
        c4 += 2.0;
        /* c2 = y - i,  c3 = i(y - i),  c4 = d + 2i,  for i even */
        a2 = c4 * a1 + c3 * a2;
        b2 = c4 * b1 + c3 * b2;

        if b2 > scalefactor {
            a1 /= scalefactor;
            b1 /= scalefactor;
            a2 /= scalefactor;
            b2 /= scalefactor
        }

        if b2 != 0.0 {
            f = a2 / b2;
            /* convergence check: relative; "absolute" for very small f : */
            if (f - of).abs() <= 2.2204460492503131e-16 * f0.max(f.abs()) {
                return f;
            }
            of = f
        }
    }

    warn!(" ** NON-convergence in pgamma()\'s pd_lower_cf() f= {}.", f);
    f
    /* should not happen ... */
}

fn pd_lower_series(lambda: f64, mut y: f64) -> f64 {
    let mut term = 1.0;
    let mut sum = 0.0;

    while y >= 1.0 && term > sum * 2.2204460492503131e-16 {
        term *= y / lambda;
        sum += term;
        y -= 1.
    }
    /* sum =  \sum_{n=0}^ oo  y*(y-1)*...*(y - n) / lambda^(n+1)
     *    =  y/lambda * (1 + \sum_{n=1}^Inf  (y-1)*...*(y-n) / lambda^n)
     *    ~  y/lambda + o(y/lambda)
     */

    if y != y.floor() {
        /*
         * The series does not converge as the terms start getting
         * bigger (besides flipping sign) for y < -lambda.
         */
        let mut f = 0.0;
        /* FIXME: in quite few cases, adding  term*f  has no effect (f too small)
         *   and is unnecessary e.g. for pgamma(4e12, 121.1) */
        f = pd_lower_cf(y, lambda + 1.0 - y);
        sum += term * f
    }

    sum
}

/// Compute the following ratio with higher accuracy that would be had
/// from doing it directly.
///
/// $ \frac{dnorm (x, 0, 1, false)}{pnorm (x, 0, 1, lower_tail, false)} $
///
/// Abramowitz & Stegun 26.2.12
fn dpnorm(mut x: f64, mut lower_tail: bool, lp: f64) -> f64 {
    /*
     * So as not to repeat a pnorm call, we expect
     *
     *  lp == pnorm (x, 0, 1, lower_tail, true)
     *
     * but use it only in the non-critical case where either x is small
     * or p==exp(lp) is close to 1.
     */
    if x < 0.0 {
        x = -x;
        lower_tail = !lower_tail
    }

    if x > 10.0 && !lower_tail {
        let mut term = 1.0 / x;
        let mut sum = term;
        let x2 = x * x;
        let mut i = 1.0;

        loop {
            term *= -i / x2;
            sum += term;
            i += 2.0;
            if !(term.abs() > 2.2204460492503131e-16 * sum) {
                break;
            }
        }

        1.0 / sum
    } else {
        let d = dnorm(x, 0.0, 1.0, false).unwrap();
        d / lp.exp()
    }
}

/// Asymptotic expansion to calculate the probability that Poisson variate
/// has value <= x.
/// Various assertions about this are made (without proof) at
/// <http://members.aol.com/iandjmsmith/PoissonApprox.htm>
fn ppois_asymp(x: f64, lambda: f64, lower_tail: bool, log: bool) -> f64 {
    static coefs_a: [f64; 8] = [
        -1e99, /* placeholder used for 1-indexing */
        2.0 / 3.0,
        -4.0 / 135.0,
        8.0 / 2835.0,
        16.0 / 8505.0,
        -8992.0 / 12629925.0,
        -334144.0 / 492567075.0,
        698752.0 / 1477701225.0,
    ];

    static coefs_b: [f64; 8] = [
        -1e99, /* placeholder */
        1.0 / 12.0,
        1.0 / 288.0,
        -139.0 / 51840.0,
        -571.0 / 2488320.0,
        163879.0 / 209018880.0,
        5246819.0 / 75246796800.0,
        -534703531.0 / 902961561600.0,
    ];

    let mut elfb = 0.0;
    let mut elfb_term = 0.0;
    let mut res12 = 0.0;
    let mut res1_term = 0.0;
    let mut res1_ig = 0.0;
    let mut res2_term = 0.0;
    let mut res2_ig = 0.0;
    let mut dfm = 0.0;
    let mut pt_ = 0.0;
    let mut s2pt = 0.0;
    let mut f = 0.0;
    let mut np = 0.0;
    let mut i = 0;

    dfm = lambda - x;
    /* If lambda is large, the distribution is highly concentrated
       about lambda.  So representation error in x or lambda can lead
       to arbitrarily large values of pt_ and hence divergence of the
       coefficients of this approximation.
    */
    pt_ = -log1pmx(dfm / x);
    s2pt = (2.0 * x * pt_).sqrt();
    if dfm < 0.0 {
        s2pt = -s2pt
    }

    res12 = 0.0;
    res1_term = x.sqrt();
    res1_ig = res1_term;
    res2_term = s2pt;
    res2_ig = res2_term;
    i = 1;
    while i < 8 {
        res12 += res1_ig * coefs_a[i];
        res12 += res2_ig * coefs_b[i];
        res1_term *= pt_ / i as f64;
        res2_term *= 2.0 * pt_ / (2 * i + 1) as f64;
        res1_ig = res1_ig / x + res1_term;
        res2_ig = res2_ig / x + res2_term;
        i += 1
    }

    elfb = x;
    elfb_term = 1.0;
    i = 1;
    while i < 8 {
        elfb += elfb_term * coefs_b[i];
        elfb_term /= x;
        i += 1
    }
    if !lower_tail {
        elfb = -elfb
    }

    f = res12 / elfb;

    np = if log {
        log_pnorm(s2pt, 0.0, 1.0, !lower_tail).unwrap()
    } else {
        pnorm(s2pt, 0.0, 1.0, !lower_tail).unwrap()
    };

    if log {
        let n_d_over_p = dpnorm(s2pt, !lower_tail, np);
        np + (f * n_d_over_p).ln_1p()
    } else {
        let nd = dnorm(s2pt, 0.0, 1.0, log).unwrap();
        np + f * nd
    }
}

pub fn pgamma_raw(x: f64, alph: f64, lower_tail: bool, log: bool) -> f64 {
    /* Here, assume that  (x,alph) are not NA  &  alph > 0 . */
    let mut res = 0.0;

    if let Some(ret) = x.p_bounds_01(0.0, f64::infinity(), lower_tail, log) {
        return ret;
    }

    if x < 1.0 {
        res = pgamma_smallx(x, alph, lower_tail, log)
    } else if x <= alph - 1.0 && x < 0.8 * (alph + 50.0) {
        /* incl. large alph compared to x */
        let sum = pd_upper_series(x, alph, log); /* = x/alph + o(x/alph) */
        let d = dpois_wrap(alph, x, log); /* x >= 1 and x fairly near alph. */
        if !lower_tail {
            res = if log {
                (d + sum).log1_exp()
            } else {
                (1.0) - d * sum
            }
        } else {
            res = if log { (sum) + d } else { (sum) * d }
        }
    } else if (alph - 1.0) < x && alph < 0.8 * (x + 50.0) {
        /* incl. large x compared to alph */
        let mut sum_0 = 0.0;
        let d_0 = dpois_wrap(alph, x, log);
        if alph < 1.0 {
            if x * 2.2204460492503131e-16 > 1.0 - alph {
                sum_0 = f64::d_1(log);
            } else {
                let f = pd_lower_cf(alph, x - (alph - 1.0)) * x / alph;
                /* = [alph/(x - alph+1) + o(alph/(x-alph+1))] * x/alph = 1 + o(1) */
                sum_0 = if log { f.ln() } else { f }
            }
        } else {
            sum_0 = pd_lower_series(x, alph - 1.0); /* = (alph-1)/x + o((alph-1)/x) */
            sum_0 = if log { sum_0.ln_1p() } else { (1.0) + sum_0 }
        }
        if !lower_tail {
            res = if log { (sum_0) + d_0 } else { (sum_0) * d_0 }
        } else {
            res = if log {
                (d_0 + sum_0).log1_exp()
            } else {
                (1.0) - d_0 * sum_0
            }
        }
    } else {
        res = ppois_asymp(alph - 1.0, x, !lower_tail, log)
    }

    /*
     * We lose a fair amount of accuracy to underflow in the cases
     * where the final result is very close to DBL_MIN.  In those
     * cases, simply redo via log space.
     */
    if !log && res < 2.2250738585072014e-308 / 2.2204460492503131e-16 {
        /* with(.Machine, double.xmin / double.eps) #|-> 1.002084e-292 */
        pgamma_raw(x, alph, lower_tail, true).exp()
    } else {
        res
    }
}

/// This function computes the distribution function for the
/// gamma distribution with shape parameter alph and scale parameter
/// scale. This is also known as the incomplete gamma function.
/// See Abramowitz and Stegun (6.5.1) for example.
pub fn pgamma<RE: Into<Real64>, P: Into<Positive64>, RA: Into<Rational64>>(
    x: RE,
    alph: P,
    scale: RA,
    lower_tail: bool,
) -> Probability64 {
    pgamma_inner(x, alph, scale, lower_tail, false).into()
}

/// This function computes the distribution function for the
/// gamma distribution with shape parameter alph and scale parameter
/// scale. This is also known as the incomplete gamma function.
/// See Abramowitz and Stegun (6.5.1) for example.
pub fn log_pgamma<RE: Into<Real64>, P: Into<Positive64>, RA: Into<Rational64>>(
    x: RE,
    alph: P,
    scale: RA,
    lower_tail: bool,
) -> LogProbability64 {
    pgamma_inner(x, alph, scale, lower_tail, true).into()
}

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

    x /= scale;
    if alph == 0.0 {
        /* limit case; useful e.g. in pnchisq() */
        return if x <= 0.0 {
            f64::dt_0(lower_tail, log)
        } else {
            f64::dt_1(lower_tail, log)
        };
    }
    pgamma_raw(x, alph, lower_tail, log)
}