r2rs-nmath 0.1.1

// Translation of nmath's rbinom
//
// 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, PositiveInteger64, Probability64, Real64};

use crate::{distribution::binom::qbinom, traits::RNG};

thread_local! {
    static C: RefCell<f64> = RefCell::new(0.0);
    static FM: RefCell<f64> = RefCell::new(0.0);
    static NPQ: RefCell<f64> = RefCell::new(0.0);
    static P1: RefCell<f64> = RefCell::new(0.0);
    static P2: RefCell<f64> = RefCell::new(0.0);
    static P3: RefCell<f64> = RefCell::new(0.0);
    static P4: RefCell<f64> = RefCell::new(0.0);
    static QN: RefCell<f64> = RefCell::new(0.0);
    static XL: RefCell<f64> = RefCell::new(0.0);
    static XLL: RefCell<f64> = RefCell::new(0.0);
    static XLR: RefCell<f64> = RefCell::new(0.0);
    static XM: RefCell<f64> = RefCell::new(0.0);
    static XR: RefCell<f64> = RefCell::new(0.0);
    static PSAVE: RefCell<f64> = RefCell::new(-1.0);
    static NSAVE: RefCell<i32> = RefCell::new(-1);
    static M: RefCell<i32> = RefCell::new(0);
}

pub fn reset_rbinom_statics() {
    C.with(|C_| *C_.borrow_mut() = 0.0);
    FM.with(|FM_| *FM_.borrow_mut() = 0.0);
    NPQ.with(|NPQ_| *NPQ_.borrow_mut() = 0.0);
    P1.with(|P1_| *P1_.borrow_mut() = 0.0);
    P2.with(|P2_| *P2_.borrow_mut() = 0.0);
    P3.with(|P3_| *P3_.borrow_mut() = 0.0);
    P4.with(|P4_| *P4_.borrow_mut() = 0.0);
    QN.with(|QN_| *QN_.borrow_mut() = 0.0);
    XL.with(|XL_| *XL_.borrow_mut() = 0.0);
    XLL.with(|XLL_| *XLL_.borrow_mut() = 0.0);
    XLR.with(|XLR_| *XLR_.borrow_mut() = 0.0);
    XM.with(|XM_| *XM_.borrow_mut() = 0.0);
    XR.with(|XR_| *XR_.borrow_mut() = 0.0);
    PSAVE.with(|PSAVE_| *PSAVE_.borrow_mut() = -1.0);
    NSAVE.with(|NSAVE_| *NSAVE_.borrow_mut() = -1);
    M.with(|M_| *M_.borrow_mut() = 0);
}

/// Random variates from the binomial distribution.
pub fn rbinom<PO: Into<PositiveInteger64>, PR: Into<Probability64>, R: RNG>(
    nin: PO,
    pp: PR,
    rng: &mut R,
) -> Real64 {
    let nin = nin.into().unwrap();
    let pp = pp.into().unwrap();

    let mut current_block: u64;

    let mut c = C.with(|C_| *C_.borrow());
    let mut fm = FM.with(|FM_| *FM_.borrow());
    let mut npq = NPQ.with(|NPQ_| *NPQ_.borrow());
    let mut p1 = P1.with(|P1_| *P1_.borrow());
    let mut p2 = P2.with(|P2_| *P2_.borrow());
    let mut p3 = P3.with(|P3_| *P3_.borrow());
    let mut p4 = P4.with(|P4_| *P4_.borrow());
    let mut qn = QN.with(|QN_| *QN_.borrow());
    let mut xl = XL.with(|XL_| *XL_.borrow());
    let mut xll = XLL.with(|XLL_| *XLL_.borrow());
    let mut xlr = XLR.with(|XLR_| *XLR_.borrow());
    let mut xm = XM.with(|XM_| *XM_.borrow());
    let mut xr = XR.with(|XR_| *XR_.borrow());
    let mut psave = PSAVE.with(|PSAVE_| *PSAVE_.borrow());
    let mut nsave = NSAVE.with(|NSAVE_| *NSAVE_.borrow());
    let mut m = M.with(|M_| *M_.borrow());

    let mut f = 0.0;
    let mut f1 = 0.0;
    let mut f2 = 0.0;
    let mut u = 0.0;
    let mut v = 0.0;
    let mut w = 0.0;
    let mut w2 = 0.0;
    let mut x = 0.0;
    let mut x1 = 0.0;
    let mut x2 = 0.0;
    let mut z = 0.0;
    let mut z2 = 0.0;
    let mut p = 0.0;
    let mut q = 0.0;
    let mut np = 0.0;
    let mut g = 0.0;
    let mut r = 0.0;
    let mut al = 0.0;
    let mut alv = 0.0;
    let mut amaxp = 0.0;
    let mut ffm = 0.0;
    let mut ynorm = 0.0;
    let mut i = 0;
    let mut ix = 0;
    let mut k = 0;
    let mut n = 0;

    r = nin;
    if !pp.is_finite() {
        // n=0, p=0, p=1 are not errors <TSL>
        return f64::nan().into();
    }

    if r == 0.0 || pp == 0.0 {
        return 0.0.into();
    }
    if pp == 1.0 {
        return r.into();
    }

    if r >= std::i32::MAX as f64 {
        /* evade integer overflow,
        and r == INT_MAX gave only even values */
        return qbinom(rng.unif_rand(), r, pp, false);
    }
    /* else */
    n = r as i32;

    p = pp.min(1.0 - pp);
    q = 1.0 - p;
    np = n as f64 * p;
    r = p / q;
    g = r * (n + 1) as f64;

    /* Setup, perform only when parameters change [using static (globals): */

    /* FIXING: Want this thread safe
       -- use as little (thread globals) as possible
    */
    if pp != psave || n != nsave {
        PSAVE.with(|PSAVE_| {
            *PSAVE_.borrow_mut() = pp;
            psave = *PSAVE_.borrow();
        });
        NSAVE.with(|NSAVE_| {
            *NSAVE_.borrow_mut() = n;
            nsave = *NSAVE_.borrow();
        });
        if np < 30.0 {
            /* inverse cdf logic for mean less than 30 */
            QN.with(|QN_| {
                *QN_.borrow_mut() = q.powi(n);
                qn = *QN_.borrow();
            });
            current_block = 5089124893069931607;
        } else {
            ffm = np + p;
            M.with(|M_| {
                *M_.borrow_mut() = ffm as i32;
                m = *M_.borrow();
            });
            FM.with(|FM_| {
                *FM_.borrow_mut() = m as f64;
                fm = *FM_.borrow();
            });
            NPQ.with(|NPQ_| {
                *NPQ_.borrow_mut() = np * q;
                npq = *NPQ_.borrow();
            });
            P1.with(|P1_| {
                *P1_.borrow_mut() = (2.195 * npq.sqrt() - 4.6 * q).trunc() + 0.5;
                p1 = *P1_.borrow();
            });
            XM.with(|XM_| {
                *XM_.borrow_mut() = fm + 0.5;
                xm = *XM_.borrow();
            });
            XL.with(|XL_| {
                *XL_.borrow_mut() = xm - p1;
                xl = *XL_.borrow();
            });
            XR.with(|XR_| {
                *XR_.borrow_mut() = xm + p1;
                xr = *XR_.borrow();
            });
            C.with(|C_| {
                *C_.borrow_mut() = 0.134 + 20.5 / (15.3 + fm);
                c = *C_.borrow();
            });
            al = (ffm - xl) / (ffm - xl * p);
            XLL.with(|XLL_| {
                *XLL_.borrow_mut() = al * (1.0 + 0.5 * al);
                xll = *XLL_.borrow();
            });
            al = (xr - ffm) / (xr * q);
            XLR.with(|XLR_| {
                *XLR_.borrow_mut() = al * (1.0 + 0.5 * al);
                xlr = *XLR_.borrow();
            });
            P2.with(|P2_| {
                *P2_.borrow_mut() = p1 * (1.0 + c + c);
                p2 = *P2_.borrow();
            });
            P3.with(|P3_| {
                *P3_.borrow_mut() = p2 + c / xll;
                p3 = *P3_.borrow();
            });
            P4.with(|P4_| {
                *P4_.borrow_mut() = p3 + c / xlr;
                p4 = *P4_.borrow();
            });
            current_block = 17769492591016358583;
        }
    } else if n == nsave {
        if np < 30.0 {
            current_block = 5089124893069931607;
        } else {
            current_block = 17769492591016358583;
        }
    } else {
        current_block = 17769492591016358583;
    }
    's_465: loop {
        match current_block {
            5089124893069931607 =>
            /*---------------------- np = n*p < 30 : ------------------------- */
            {
                ix = 0;
                f = qn;
                u = rng.unif_rand();
                loop {
                    if u < f {
                        break 's_465;
                    }
                    if ix > 110 {
                        current_block = 5089124893069931607;
                        break;
                    }
                    u -= f;
                    ix += 1;
                    f *= g / ix as f64 - r
                }
            }
            _ =>
            /*-------------------------- np = n*p >= 30 : ------------------- */
            {
                u = rng.unif_rand() * p4;
                v = rng.unif_rand();
                /* triangular region */
                if u <= p1 {
                    ix = (xm - p1 * v + u) as i32;
                    break;
                } else {
                    /* parallelogram region */
                    if u <= p2 {
                        x = xl + (u - p1) / c;
                        v = v * c + 1.0 - (xm - x).abs() / p1;
                        if v > 1.0 || v <= 0.0 {
                            current_block = 17769492591016358583;
                            continue;
                        }
                        ix = x as i32
                    } else if u > p3 {
                        /* right tail */
                        ix = (xr - v.ln() / xlr) as i32;
                        if ix > n {
                            current_block = 17769492591016358583;
                            continue;
                        }
                        v = v * (u - p3) * xlr
                    } else {
                        /* left tail */
                        ix = (xl + v.ln() / xll) as i32;
                        if ix < 0 {
                            current_block = 17769492591016358583;
                            continue;
                        }
                        v = v * (u - p2) * xll
                    }
                    /* determine appropriate way to perform accept/reject test */
                    k = (ix - m).abs();
                    if k <= 20 || k as f64 >= npq / 2.0 - 1.0 {
                        /* explicit evaluation */
                        f = 1.0;
                        if m < ix {
                            i = m + 1;
                            while i <= ix {
                                f *= g / i as f64 - r;
                                i += 1
                            }
                        } else if m != ix {
                            i = ix + 1;
                            while i <= m {
                                f /= g / i as f64 - r;
                                i += 1
                            }
                        }
                        if v <= f {
                            break;
                        } else {
                            current_block = 17769492591016358583;
                        }
                    } else {
                        /* squeezing using upper and lower bounds on f(x).ln() */
                        amaxp = (k as f64 / npq)
                            * ((k as f64 * (k as f64 / 3.0 + 0.625) + 0.1666666666666) / npq + 0.5);
                        ynorm = -k as f64 * k as f64 / (2.0 * npq);
                        alv = v.ln();
                        if alv < ynorm - amaxp {
                            break;
                        }
                        if !(alv <= ynorm + amaxp) {
                            current_block = 17769492591016358583;
                            continue;
                        }
                        /* stirling's formula to machine accuracy */
                        /* for the final acceptance/rejection test */
                        x1 = ix as f64 + 1.0;
                        f1 = fm + 1.0;
                        z = (n + 1) as f64 - fm;
                        w = (n as f64 - ix as f64) + 1.0;
                        z2 = z * z;
                        x2 = x1 * x1;
                        f2 = f1 * f1;
                        w2 = w * w;
                        if alv
                            <= xm * (f1 / x1).ln()
                                + ((n as f64 - m as f64) + 0.5) * (z / w).ln()
                                + (ix as f64 - m as f64) * (w * p / (x1 * q)).ln()
                                + (13860.0 - (462.0 - (132.0 - (99.0 - 140.0 / f2) / f2) / f2) / f2)
                                    / f1
                                    / 166320.0
                                + (13860.0 - (462.0 - (132.0 - (99.0 - 140.0 / z2) / z2) / z2) / z2)
                                    / z
                                    / 166320.0
                                + (13860.0 - (462.0 - (132.0 - (99.0 - 140.0 / x2) / x2) / x2) / x2)
                                    / x1
                                    / 166320.0
                                + (13860.0 - (462.0 - (132.0 - (99.0 - 140.0 / w2) / w2) / w2) / w2)
                                    / w
                                    / 166320.0
                        {
                            break;
                        } else {
                            current_block = 17769492591016358583;
                        }
                    }
                }
            }
        }
    }
    if psave > 0.5 {
        ix = n - ix
    }

    ix.into()
}