r2rs-nmath 0.1.1

// Translation of nmath's rgamma
//
// 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::{cell::RefCell, thread_local};

use num_traits::Float;
use strafe_type::{FloatConstraint, Positive64, Rational64, Real64};

use crate::{rng::NMathRNG, traits::RNG};

thread_local! {
    static AA: RefCell<f64> = RefCell::new(0.0); /* no. 1 (step 1) */
    static AAA: RefCell<f64> = RefCell::new(0.0); /* no. 2 (step 4) */
    static S: RefCell<f64> = RefCell::new(0.0);
    static S2: RefCell<f64> = RefCell::new(0.0);
    static D: RefCell<f64> = RefCell::new(0.0);
    static Q0: RefCell<f64> = RefCell::new(0.0);
    static B: RefCell<f64> = RefCell::new(0.0);
    static SI: RefCell<f64> = RefCell::new(0.0);
    static C: RefCell<f64> = RefCell::new(0.0);
}

pub fn reset_rgamma_statics() {
    AA.with(|AA_| *AA_.borrow_mut() = 0.0);
    AAA.with(|AAA_| *AAA_.borrow_mut() = 0.0);
    S.with(|S_| *S_.borrow_mut() = 0.0);
    S2.with(|S2_| *S2_.borrow_mut() = 0.0);
    D.with(|D_| *D_.borrow_mut() = 0.0);
    Q0.with(|Q0_| *Q0_.borrow_mut() = 0.0);
    B.with(|B_| *B_.borrow_mut() = 0.0);
    SI.with(|SI_| *SI_.borrow_mut() = 0.0);
    C.with(|C_| *C_.borrow_mut() = 0.0);
}

/// Random variates from the gamma distribution.
///
/// Input : a     = 'shape' = alpha,
/// scale = 'scale' = 1/rate of the gamma distribution w/ mean E[.] = a * scale
/// Output: a variate from the gamma(a, scale)-distribution
pub fn rgamma<P: Into<Positive64>, RA: Into<Rational64>, R: RNG>(
    a: P,
    scale: RA,
    rng: &mut R,
) -> Real64 {
    let a = a.into().unwrap();
    let scale = scale.into().unwrap();

    /* Constants : */
    static sqrt32: f64 = 5.656854; /* exp(-1) = 1/e */
    static exp_m1: f64 = 0.36787944117144232159;

    /* Coefficients q[k] - for q0 = sum(q[k]*a^(-k))
     * Coefficients a[k] - for q = q0+(t*t/2)*sum(a[k]*v^k)
     * Coefficients e[k] - for exp(q)-1 = sum(e[k]*q^k)
     */
    static q1: f64 = 0.04166669;
    static q2: f64 = 0.02083148;
    static q3: f64 = 0.00801191;
    static q4: f64 = 0.00144121;
    static q5: f64 = -7.388e-5;
    static q6: f64 = 2.4511e-4;
    static q7: f64 = 2.424e-4;

    static a1: f64 = 0.3333333;
    static a2: f64 = -0.250003;
    static a3: f64 = 0.2000062;
    static a4: f64 = -0.1662921;
    static a5: f64 = 0.1423657;
    static a6: f64 = -0.1367177;
    static a7: f64 = 0.1233795;

    let mut aa = AA.with(|AA_| *AA_.borrow()); /* no. 1 (step 1) */
    let mut aaa = AAA.with(|AAA_| *AAA_.borrow()); /* no. 2 (step 4) */
    let mut s = S.with(|S_| *S_.borrow());
    let mut s2 = S2.with(|S2_| *S2_.borrow());
    let mut d = D.with(|D_| *D_.borrow());
    let mut q0 = Q0.with(|Q0_| *Q0_.borrow());
    let mut b = B.with(|B_| *B_.borrow());
    let mut si = SI.with(|SI_| *SI_.borrow());
    let mut c = C.with(|C_| *C_.borrow());

    let mut e = 0.0;
    let mut p = 0.0;
    let mut q = 0.0;
    let mut r = 0.0;
    let mut t = 0.0;
    let mut u = 0.0;
    let mut v = 0.0;
    let mut w = 0.0;
    let mut x = 0.0;

    let mut ret_val = 0.0;

    if a == 0.0 {
        return 0.0.into();
    }
    if !a.is_finite() || !scale.is_finite() {
        return f64::infinity().into();
    }

    if a < 1.0 {
        /* GS algorithm for parameters a < 1 */
        e = 1.0 + exp_m1 * a;
        loop {
            p = e * rng.unif_rand();
            if p >= 1.0 {
                x = -((e - p) / a).ln();
                if rng.exp_rand() >= (1.0 - a) * x.ln() {
                    break;
                }
            } else {
                x = (p.ln() / a).exp();
                if rng.exp_rand() >= x {
                    break;
                }
            }
        }
        return (scale * x).into();
    }

    /* --- a >= 1 : GD algorithm --- */
    /* Step 1: Recalculations of s2, s, d if a has changed */
    if a != aa {
        AA.with(|AA_| {
            *AA_.borrow_mut() = a;
            aa = *AA_.borrow();
        });
        S2.with(|S2_| {
            *S2_.borrow_mut() = a - 0.5;
            s2 = *S2_.borrow();
        });
        S.with(|S_| {
            *S_.borrow_mut() = s2.sqrt();
            s = *S_.borrow();
        });
        D.with(|D_| {
            *D_.borrow_mut() = sqrt32 - s * 12.0;
            d = *D_.borrow();
        });
    }

    /* Step 2: t = standard normal deviate,
    x = (s,1/2) -normal deviate. */
    /* immediate acceptance (i) */
    t = rng.norm_rand();
    x = s + 0.5 * t;
    ret_val = x * x;
    if t >= 0.0 {
        return (scale * ret_val).into();
    }

    /* Step 3: u = 0,1 - uniform sample. squeeze acceptance (s) */
    u = rng.unif_rand();
    if d * u <= t * t * t {
        return (scale * ret_val).into();
    }

    /* Step 4: recalculations of q0, b, si, c if necessary */
    if a != aaa {
        AAA.with(|AAA_| {
            *AAA_.borrow_mut() = a;
            aaa = *AAA_.borrow();
        });
        r = 1.0 / a;
        Q0.with(|Q0_| {
            *Q0_.borrow_mut() =
                ((((((q7 * r + q6) * r + q5) * r + q4) * r + q3) * r + q2) * r + q1) * r;
            q0 = *Q0_.borrow();
        });

        /* Approximation depending on size of parameter a */
        /* The constants in the expressions for b, si and c */
        /* were established by numerical experiments */

        if a <= 3.686 {
            B.with(|B_| {
                *B_.borrow_mut() = 0.463 + s + 0.178 * s2;
                b = *B_.borrow();
            });
            SI.with(|SI_| {
                *SI_.borrow_mut() = 1.235;
                si = *SI_.borrow();
            });
            C.with(|C_| {
                *C_.borrow_mut() = 0.195 / s - 0.079 + 0.16 * s;
                c = *C_.borrow();
            });
        } else if a <= 13.022 {
            B.with(|B_| {
                *B_.borrow_mut() = 1.654 + 0.0076 * s2;
                b = *B_.borrow();
            });
            SI.with(|SI_| {
                *SI_.borrow_mut() = 1.68 / s + 0.275;
                si = *SI_.borrow();
            });
            C.with(|C_| {
                *C_.borrow_mut() = 0.062 / s + 0.024;
                c = *C_.borrow();
            });
        } else {
            B.with(|B_| {
                *B_.borrow_mut() = 1.77;
                b = *B_.borrow();
            });
            SI.with(|SI_| {
                *SI_.borrow_mut() = 0.75;
                si = *SI_.borrow();
            });
            C.with(|C_| {
                *C_.borrow_mut() = 0.1515 / s;
                c = *C_.borrow();
            });
        }
    }

    /* Step 5: no quotient test if x not positive */
    if x > 0.0 {
        /* Step 6: calculation of v and quotient q */
        v = t / (s + s);
        if v.abs() <= 0.25 {
            q = q0
                + 0.5
                    * t
                    * t
                    * ((((((a7 * v + a6) * v + a5) * v + a4) * v + a3) * v + a2) * v + a1)
                    * v
        } else {
            q = q0 - s * t + 0.25 * t * t + (s2 + s2) * (1.0 + v).ln()
        }

        /* Step 7: quotient acceptance (q) */
        if (1.0 - u).ln() <= q {
            return (scale * ret_val).into();
        }
    }

    /* Step 8: e = standard exponential deviate
     * u =  0,1 -uniform deviate
     * t = (b,si)-double exponential (laplace) sample
     * repeat .. until  `t' is accepted */
    loop {
        e = rng.exp_rand();
        u = rng.unif_rand();
        u = u + u - 1.0;
        if u < 0.0 {
            t = b - si * e
        } else {
            t = b + si * e
        }

        /* Step  9:  rejection if t < tau(1) = -0.71874483771719 */
        if t >= -0.71874483771719 {
            /* Step 10:  calculation of v and quotient q */
            v = t / (s + s);
            if v.abs() <= 0.25 {
                q = q0
                    + 0.5
                        * t
                        * t
                        * ((((((a7 * v + a6) * v + a5) * v + a4) * v + a3) * v + a2) * v + a1)
                        * v
            } else {
                q = q0 - s * t + 0.25 * t * t + (s2 + s2) * (1.0 + v).ln()
            }

            /* Step 11:  hat acceptance (h) */
            /* (if q not positive go to step 8) */
            if q > 0.0 {
                w = q.exp_m1();
                /*  ^^^^^ original code had approximation with rel.err < 2e-7 */
                /* if t is rejected sample again at step 8 */
                if c * u.abs() <= w * (e - 0.5 * t * t).exp() {
                    break;
                }
            }
        }
    }
    x = s + 0.5 * t;
    return (scale * x * x).into();
}