r2rs-nmath 0.1.1

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

use crate::{
    distribution::{binom::rbinom, hyper::qhyper},
    traits::RNG,
};

thread_local! {
    static KS: RefCell<usize> = RefCell::new(1); // <- not allowing larger integer par
    static N1S: RefCell<usize> = RefCell::new(1);
    static N2S: RefCell<usize> = RefCell::new(1);
    static M: RefCell<usize> = RefCell::new(0);
    static MINJX: RefCell<usize> = RefCell::new(0);
    static MAXJX: RefCell<usize> = RefCell::new(0);
    static K: RefCell<usize> = RefCell::new(0);
    static N1: RefCell<usize> = RefCell::new(0);
    static N2: RefCell<usize> = RefCell::new(0);
    static N: RefCell<f64> = RefCell::new(0.0);
    // II :
    static W: RefCell<f64> = RefCell::new(0.0);
    // III:
    static A: RefCell<f64> = RefCell::new(0.0);
    static D: RefCell<f64> = RefCell::new(0.0);
    static S: RefCell<f64> = RefCell::new(0.0);
    static XL: RefCell<f64> = RefCell::new(0.0);
    static XR: RefCell<f64> = RefCell::new(0.0);
    static KL: RefCell<f64> = RefCell::new(0.0);
    static KR: RefCell<f64> = RefCell::new(0.0);
    static LAMDL: RefCell<f64> = RefCell::new(0.0);
    static LAMDR: 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);
}

pub fn reset_rhyper_statics() {
    KS.with(|KS_| *KS_.borrow_mut() = 1);
    N1S.with(|N1S_| *N1S_.borrow_mut() = 1);
    N2S.with(|N2S_| *N2S_.borrow_mut() = 1);
    M.with(|M_| *M_.borrow_mut() = 0);
    MINJX.with(|MINJX_| *MINJX_.borrow_mut() = 0);
    MAXJX.with(|MAXJX_| *MAXJX_.borrow_mut() = 0);
    K.with(|K_| *K_.borrow_mut() = 0);
    N1.with(|N1_| *N1_.borrow_mut() = 0);
    N2.with(|N2_| *N2_.borrow_mut() = 0);
    N.with(|N_| *N_.borrow_mut() = 0.0);
    W.with(|W_| *W_.borrow_mut() = 0.0);
    A.with(|A_| *A_.borrow_mut() = 0.0);
    D.with(|D_| *D_.borrow_mut() = 0.0);
    S.with(|S_| *S_.borrow_mut() = 0.0);
    XL.with(|XL_| *XL_.borrow_mut() = 0.0);
    XR.with(|XR_| *XR_.borrow_mut() = 0.0);
    KL.with(|KL_| *KL_.borrow_mut() = 0.0);
    KR.with(|KR_| *KR_.borrow_mut() = 0.0);
    LAMDL.with(|LAMDL_| *LAMDL_.borrow_mut() = 0.0);
    LAMDR.with(|LAMDR_| *LAMDR_.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);
}

/// afc(i) :=  ln( i! ) \[logarithm of the factorial i\] = {R:} lgamma(i + 1) = {C:} lgammafn(i + 1)
fn afc(i: usize) -> f64 {
    // If (i > 7), use Stirling's approximation, otherwise use table lookup.
    static al: [f64; 8] = [
        0.0,                                /*ln(0!)=ln(1)*/
        0.0,                                /*ln(1!)=ln(1)*/
        std::f64::consts::LN_2,             /*ln(2) */
        1.79175946922805500081247735838070, /*ln(6) */
        3.17805383034794561964694160129705, /*ln(24)*/
        4.78749174278204599424770093452324,
        6.57925121201010099506017829290394,
        8.52516136106541430016553103634712,
        /* 10.60460290274525022841722740072165, approx. value below =
           10.6046028788027; rel.error = 2.26 10^{-9}

          FIXME: Use constants and if(n > ..) decisions from ./stirlerr.c
          -----  will be even *faster* for n > 500 (or so)
        */
    ];

    if i <= 7 {
        return al[i];
    }
    // else i >= 8 :
    let di = i;
    let i2 = di * di;
    (di as f64 + 0.5) * (di as f64).ln() - di as f64
        + strafe_consts::LN_SQRT_2TPI
        + (0.0833333333333333 - 0.00277777777777778 / i2 as f64) / di as f64
}

/// Random variates from the hypergeometric distribution.
/// Returns the number of white balls drawn when kk balls
/// are drawn at random from an urn containing nn1 white
/// and nn2 black balls.
///
/// rhyper(NR, NB, n) -- NR 'red', NB 'blue', n drawn, how many are 'red'
pub fn rhyper<
    P1: Into<PositiveInteger64>,
    P2: Into<PositiveInteger64>,
    P3: Into<PositiveInteger64>,
    R: RNG,
>(
    nn1in: P1,
    nn2in: P2,
    kkin: P3,
    rng: &mut R,
) -> Real64 {
    let mut nn1in = nn1in.into().unwrap();
    let mut nn2in = nn2in.into().unwrap();
    let mut kkin = kkin.into().unwrap();

    /* extern double afc(int); */

    let mut nn1 = 0;
    let mut nn2 = 0;
    let mut kk = 0;
    let mut ix = 0; // return value (coerced to double at the very end)
    let mut setup1 = false;
    let mut setup2 = false;

    let mut ks = KS.with(|KS_| *KS_.borrow());
    let mut n1s = N1S.with(|N1S_| *N1S_.borrow());
    let mut n2s = N2S.with(|N2S_| *N2S_.borrow());
    let mut m = M.with(|M_| *M_.borrow());
    let mut minjx = MINJX.with(|MINJX_| *MINJX_.borrow());
    let mut maxjx = MAXJX.with(|MAXJX_| *MAXJX_.borrow());
    // <- not allowing larger integer par
    let mut k = K.with(|K_| *K_.borrow());
    let mut n1 = N1.with(|N1_| *N1_.borrow());
    let mut n2 = N2.with(|N2_| *N2_.borrow());
    let mut n = N.with(|TN_| *TN_.borrow());

    // II :
    let mut w = W.with(|W_| *W_.borrow());
    // III:
    let mut a = A.with(|A_| *A_.borrow());
    let mut d = D.with(|D_| *D_.borrow());
    let mut s = S.with(|S_| *S_.borrow());
    let mut xl = XL.with(|XL_| *XL_.borrow());
    let mut xr = XR.with(|XR_| *XR_.borrow());
    let mut kl = KL.with(|KL_| *KL_.borrow());
    let mut kr = KR.with(|KR_| *KR_.borrow());
    let mut lamdl = LAMDL.with(|LAMDL_| *LAMDL_.borrow());
    let mut lamdr = LAMDR.with(|LAMDR_| *LAMDR_.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());

    /* check parameter validity */

    if !nn1in.is_finite() || !nn2in.is_finite() || !kkin.is_finite() {
        return f64::nan().into();
    }

    nn1in = nn1in.round();
    nn2in = nn2in.round();
    kkin = kkin.round();

    if nn1in < 0.0 || nn2in < 0.0 || kkin < 0.0 || kkin > nn1in + nn2in {
        return f64::nan().into();
    }
    if nn1in >= std::i32::MAX as f64
        || nn2in >= std::i32::MAX as f64
        || kkin >= std::i32::MAX as f64
    {
        /* large n -- evade integer overflow (and inappropriate algorithms)
        -------- */
        // FIXME: Much faster to give rbinom() approx when appropriate; -> see Kuensch(1989)
        // Johnson, Kotz,.. p.258 (top) mention the *four* different binomial approximations
        if kkin == 1.0 {
            // Bernoulli
            return rbinom(kkin, nn1in / (nn1in + nn2in), rng);
        }
        // Slow, but safe: return  F^{-1}(U)  where F(.) = phyper(.) and  U ~ U[0,1]
        return qhyper(
            rng.unif_rand(),
            nn1in,
            nn2in,
            kkin,
            false,
            // lower_tail=FALSE: a thinko, is still "correct" as equiv. to  U <--> 1-U
        );
    }
    nn1 = nn1in as usize;
    nn2 = nn2in as usize;
    kk = kkin as usize;

    /* if new parameter values, initialize */
    if nn1 != n1s || nn2 != n2s {
        // n1 | n2 is changed: setup all
        setup1 = true;
        setup2 = true
    } else if kk != ks {
        // n1 & n2 are unchanged: setup 'k' only
        setup1 = false;
        setup2 = true
    } else {
        // all three unchanged ==> no setup
        setup1 = false;
        setup2 = false
    }
    if setup1 {
        // n1 & n2
        N1S.with(|N1S_| {
            // save
            *N1S_.borrow_mut() = nn1;
            n1s = *N1S_.borrow();
        });
        N2S.with(|N2S_| {
            // save
            *N2S_.borrow_mut() = nn2;
            n2s = *N2S_.borrow();
        });
        N.with(|N_| {
            // avoid int overflow
            *N_.borrow_mut() = (nn1 + nn2) as f64;
            n = *N_.borrow();
        });
        if nn1 <= nn2 {
            N1.with(|N1_| {
                *N1_.borrow_mut() = nn1;
                n1 = *N1_.borrow();
            });
            N2.with(|N2_| {
                *N2_.borrow_mut() = nn2;
                n2 = *N2_.borrow();
            });
        } else {
            // nn2 < nn1
            N1.with(|N1_| {
                *N1_.borrow_mut() = nn2;
                n1 = *N1_.borrow();
            });
            N2.with(|N2_| {
                *N2_.borrow_mut() = nn1;
                n2 = *N2_.borrow();
            });
        }
        // now have n1 <= n2
    }
    if setup2 {
        // k
        KS.with(|KS_| {
            // save
            *KS_.borrow_mut() = kk;
            ks = *KS_.borrow();
        });
        if (kk + kk) as f64 >= n {
            // this could overflow
            K.with(|K_| {
                *K_.borrow_mut() = n as usize - kk;
                k = *K_.borrow();
            });
        } else {
            K.with(|K_| {
                *K_.borrow_mut() = kk;
                k = *K_.borrow();
            });
        }
    }
    if setup1 || setup2 {
        M.with(|M_| {
            // m := floor(adjusted mean E[.])
            *M_.borrow_mut() = ((k as f64 + 1.0) * (n1 as f64 + 1.0) / (n + 2.0)) as usize;
            m = *M_.borrow();
        });
        MINJX.with(|MINJX_| {
            *MINJX_.borrow_mut() = (k as i32 - n2 as i32).max(0) as usize;
            minjx = *MINJX_.borrow();
        });
        MAXJX.with(|MAXJX_| {
            *MAXJX_.borrow_mut() = n1.min(k);
            maxjx = *MAXJX_.borrow();
        });
    }
    /* generate random variate --- Three basic cases */

    if minjx == maxjx {
        /* I: degenerate distribution ---------------- */
        ix = maxjx;
        // return appropriate variate
        if (kk + kk) as f64 >= n {
            if nn1 > nn2 {
                ix = kk - nn2 + ix
            } else {
                ix = nn1 - ix
            }
        } else if nn1 > nn2 {
            ix = kk - ix
        }
        return ix.into();
    } else if m - minjx < 10 {
        static scale: f64 = 1e25; /* III : H2PE Algorithm --------------------------------------- */
        static con: f64 = 57.5646273248511421;
        // II: (Scaled) algorithm HIN (inverse transformation) ----
        if setup1 || setup2 {
            // scaling factor against (early) underflow
            // 25*10.ln() = scale.ln() { <==> con.exp() == scale }
            let mut lw = 0.0; // w.ln();  w = lw.exp() * scale = (lw + scale.ln()).exp() = (lw + con).exp()
            if k < n2 {
                lw = afc(n2) + afc(n1 + n2 - k) - afc(n2 - k) - afc(n1 + n2)
            } else {
                lw = afc(n1) + afc(k) - afc(k - n2) - afc(n1 + n2)
            }
            W.with(|W_| {
                *W_.borrow_mut() = (lw + con).exp();
                w = *W_.borrow();
            });
        }
        let mut p = 0.0;
        let mut u = 0.0;
        'c_5406: loop {
            p = w;
            ix = minjx;
            u = rng.unif_rand() * scale;
            loop {
                if !(u > p) {
                    break 'c_5406;
                }
                u -= p;
                p *= (n1 as f64 - ix as f64) * (k as f64 - ix as f64);
                ix += 1;
                p = p / ix as f64 / (n2 as f64 - k as f64 + ix as f64);
                if ix > maxjx {
                    break;
                }
                // FIXME  if(p == 0.0)  we also "have lost"  => goto L10
            }
        }
    } else {
        let mut u_0 = 0.0;
        let mut v = 0.0;
        if setup1 || setup2 {
            S.with(|S_| {
                *S_.borrow_mut() =
                    ((n - k as f64) * k as f64 * n1 as f64 * n2 as f64 / (n - 1.0) / n / n).sqrt();
                s = *S_.borrow();
            });
            /* remark: d is defined in reference without int. */
            /* the truncation centers the cell boundaries at 0.5 */
            D.with(|D_| {
                *D_.borrow_mut() = (1.5 * s).floor() + 0.5;
                d = *D_.borrow();
            });
            XL.with(|XL_| {
                *XL_.borrow_mut() = m as f64 - d as f64 + 0.5;
                xl = *XL_.borrow();
            });
            XR.with(|XR_| {
                *XR_.borrow_mut() = m as f64 + d as f64 + 0.5;
                xr = *XR_.borrow();
            });
            A.with(|A_| {
                *A_.borrow_mut() = afc(m) + afc(n1 - m) + afc(k - m) + afc(n2 - k + m);
                a = *A_.borrow();
            });
            KL.with(|KL_| {
                *KL_.borrow_mut() = (a
                    - afc(xl as usize)
                    - afc(n1 - xl as usize)
                    - afc(k - xl as usize)
                    - afc((n2 - k) + xl as usize))
                .exp();
                kl = *KL_.borrow();
            });
            KR.with(|KR_| {
                *KR_.borrow_mut() = (a
                    - afc(xr as usize - 1)
                    - afc(n1 - xr as usize + 1)
                    - afc(k - xr as usize + 1)
                    - afc((n2 - k) + xr as usize - 1))
                .exp();
                kr = *KR_.borrow();
            });
            LAMDL.with(|LAMDL_| {
                *LAMDL_.borrow_mut() = -(xl * ((n2 as f64 - k as f64) + xl)
                    / (n1 as f64 - xl + 1.0)
                    / (k as f64 - xl + 1.0))
                    .ln();
                lamdl = *LAMDL_.borrow();
            });
            LAMDR.with(|LAMDR_| {
                *LAMDR_.borrow_mut() = -((n1 as f64 - xr + 1.0) * (k as f64 - xr + 1.0)
                    / xr
                    / ((n2 as f64 - k as f64) + xr))
                    .ln();
                lamdr = *LAMDR_.borrow();
            });
            P1.with(|P1_| {
                *P1_.borrow_mut() = d + d;
                p1 = *P1_.borrow();
            });
            P2.with(|P2_| {
                *P2_.borrow_mut() = p1 + kl / lamdl;
                p2 = *P2_.borrow();
            });
            P3.with(|P3_| {
                *P3_.borrow_mut() = p2 + kr / lamdr;
                p3 = *P3_.borrow();
            });
        }
        let mut n_uv = 0;
        loop {
            u_0 = rng.unif_rand() * p3;
            v = rng.unif_rand();
            n_uv += 1;
            if n_uv >= 10000 {
                warn!("rhyper() branch III: giving up after {} rejections", n_uv);

                return f64::nan().into();
            }
            if u_0 < p1 {
                /* rectangular region */
                ix = (xl + u_0) as usize
            } else if u_0 <= p2 {
                ix = (xl + v.ln() / lamdl) as usize; /* right tail */
                if ix < minjx {
                    continue;
                }
                v = v * (u_0 - p1) * lamdl
            } else {
                ix = (xr - v.ln() / lamdr) as usize;
                if ix > maxjx {
                    continue;
                }
                v = v * (u_0 - p2) * lamdr
            }
            /* left tail */
            /* acceptance/rejection test */
            let mut reject = true; // else
            if m < 100 || ix <= 50 {
                /* explicit evaluation */
                /* The original algorithm (and TOMS 668) have
                f = f * i * (n2 - k + i) / (n1 - i) / (k - i);
                in the (m > ix) case, but the definition of the
                recurrence relation on p134 shows that the +1 is
                needed. */
                let mut i = 0;
                let mut f = 1.0;
                if m < ix {
                    i = m + 1;
                    while i <= ix {
                        f = f * (n1 - i + 1) as f64 * (k - i + 1) as f64
                            / (n2 - k + i) as f64
                            / i as f64;
                        i += 1
                    }
                } else if m > ix {
                    i = ix + 1;
                    while i <= m {
                        f = f * i as f64 * (n2 - k + i) as f64
                            / (n1 - i + 1) as f64
                            / (k - i + 1) as f64;
                        i += 1
                    }
                }
                if v <= f {
                    reject = false
                }
            } else {
                static deltal: f64 = 0.0078;
                static deltau: f64 = 0.0034;
                let mut e = 0.0;
                let mut g = 0.0;
                let mut r = 0.0;
                let mut t = 0.0;
                let mut y = 0.0;
                let mut de = 0.0;
                let mut dg = 0.0;
                let mut dr = 0.0;
                let mut ds = 0.0;
                let mut dt = 0.0;
                let mut gl = 0.0;
                let mut gu = 0.0;
                let mut nk = 0.0;
                let mut nm = 0.0;
                let mut ub = 0.0;
                let mut xk = 0.0;
                let mut xm = 0.0;
                let mut xn = 0.0;
                let mut y1 = 0.0;
                let mut ym = 0.0;
                let mut yn = 0.0;
                let mut yk = 0.0;
                let mut alv = 0.0;
                /* squeeze using upper and lower bounds */
                y = ix as f64;
                y1 = y + 1.0;
                ym = y - m as f64;
                yn = n1 as f64 - y + 1.0;
                yk = k as f64 - y + 1.0;
                nk = (n2 - k) as f64 + y1;
                r = -ym / y1;
                S.with(|S_| {
                    *S_.borrow_mut() = ym / yn;
                    s = *S_.borrow();
                });
                t = ym / yk;
                e = -ym / nk;
                g = yn * yk / (y1 * nk) - 1.0;
                dg = 1.0;
                if g < 0.0 {
                    dg = 1.0 + g
                }
                gu = g * (1.0 + g * (-0.5 + g / 3.0));
                gl = gu - 0.25 * (g * g * g * g) / dg;
                xm = m as f64 + 0.5;
                xn = (n1 - m) as f64 + 0.5;
                xk = (k - m) as f64 + 0.5;
                nm = (n2 - k) as f64 + xm;
                ub = y * gu - m as f64 * gl
                    + deltau
                    + xm * r * (1.0 + r * (-0.5 + r / 3.0))
                    + xn * s * (1.0 + s * (-0.5 + s / 3.0))
                    + xk * t * (1.0 + t * (-0.5 + t / 3.0))
                    + nm * e * (1.0 + e * (-0.5 + e / 3.0));
                /* test against upper bound */
                alv = v.ln();
                if alv > ub {
                    reject = true
                } else {
                    /* test against lower bound */
                    dr = xm * (r * r * r * r);
                    if r < 0.0 {
                        dr /= 1.0 + r
                    }
                    ds = xn * (s * s * s * s);
                    if s < 0.0 {
                        ds /= 1.0 + s
                    }
                    dt = xk * (t * t * t * t);
                    if t < 0.0 {
                        dt /= 1.0 + t
                    }
                    de = nm * (e * e * e * e);
                    if e < 0.0 {
                        de /= 1.0 + e
                    }
                    if alv < ub - 0.25 * (dr + ds + dt + de) + (y + m as f64) * (gl - gu) - deltal
                        || alv <= a - afc(ix) - afc(n1 - ix) - afc(k - ix) - afc(n2 - k + ix)
                    {
                        reject = false
                    } else {
                        reject = true
                    }
                }
            }
            if !(reject) {
                break;
            }
        }
    }
    /* * Stirling's formula to machine accuracy
     */
    /* return appropriate variate */
    if (kk + kk) as f64 >= n {
        if nn1 > nn2 {
            ix = kk - nn2 + ix
        } else {
            ix = nn1 - ix
        }
    } else if nn1 > nn2 {
        ix = kk - ix
    }
    ix.into()
}