r2rs-nmath 0.1.1

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

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

thread_local! {
    static L: RefCell<usize> = RefCell::new(0); /* integer "w/o overflow" */
    static M: RefCell<i32> = RefCell::new(0); /*, muold  = RefCell::new(0.0)*/
    static B1: RefCell<f64> = RefCell::new(0.0);
    static B2: RefCell<f64> = RefCell::new(0.0);
    static C: RefCell<f64> = RefCell::new(0.0);
    static C0: RefCell<f64> = RefCell::new(0.0);
    static C1: RefCell<f64> = RefCell::new(0.0);
    static C2: RefCell<f64> = RefCell::new(0.0);
    static C3: RefCell<f64> = RefCell::new(0.0);
    static PP: RefCell<[f64; 36]> = RefCell::new([0.0; 36]);
    static P0: RefCell<f64> = RefCell::new(0.0);
    static P: RefCell<f64> = RefCell::new(0.0);
    static Q: RefCell<f64> = RefCell::new(0.0);
    static S: RefCell<f64> = RefCell::new(0.0);
    static D: RefCell<f64> = RefCell::new(0.0);
    static OMEGA: RefCell<f64> = RefCell::new(0.0);
    static BIG_L: RefCell<f64> = RefCell::new(0.0);
    static MUPREV: RefCell<f64> = RefCell::new(0.0);
    static MUPREV2: RefCell<f64> = RefCell::new(0.0);
}

pub fn reset_rpois_statics() {
    L.with(|L_| *L_.borrow_mut() = 0);
    M.with(|M_| *M_.borrow_mut() = 0);
    B1.with(|B1_| *B1_.borrow_mut() = 0.0);
    B2.with(|B2_| *B2_.borrow_mut() = 0.0);
    C.with(|C_| *C_.borrow_mut() = 0.0);
    C0.with(|C0_| *C0_.borrow_mut() = 0.0);
    C1.with(|C1_| *C1_.borrow_mut() = 0.0);
    C2.with(|C2_| *C2_.borrow_mut() = 0.0);
    C3.with(|C3_| *C3_.borrow_mut() = 0.0);
    PP.with(|PP_| *PP_.borrow_mut() = [0.0; 36]);
    P0.with(|P0_| *P0_.borrow_mut() = 0.0);
    P.with(|P_| *P_.borrow_mut() = 0.0);
    Q.with(|Q_| *Q_.borrow_mut() = 0.0);
    S.with(|S_| *S_.borrow_mut() = 0.0);
    D.with(|D_| *D_.borrow_mut() = 0.0);
    OMEGA.with(|OMEGA_| *OMEGA_.borrow_mut() = 0.0);
    BIG_L.with(|BIG_L_| *BIG_L_.borrow_mut() = 0.0);
    MUPREV.with(|MUPREV_| *MUPREV_.borrow_mut() = 0.0);
    MUPREV2.with(|MUPREV2_| *MUPREV2_.borrow_mut() = 0.0);
}

static mut f: usize = 0;

pub fn rpois<P: Into<Positive64>, R: RNG>(mu: P, rng: &mut R) -> Real64 {
    unsafe {
        f += 1;
    }

    let mu = mu.into().unwrap();

    let mut current_block: u64;
    /*, MAYBE */
    /* Factorial Table (0:9)! */
    static fact: [f64; 10] = [
        1.0, 1.0, 2.0, 6.0, 24.0, 120.0, 720.0, 5040.0, 40320.0, 362880.0,
    ];

    static a0: f64 = -0.5;
    static a1: f64 = 0.3333333;
    static a2: f64 = -0.2500068;
    static a3: f64 = 0.2000118;
    static a4: f64 = -0.1661269;
    static a5: f64 = 0.1421878;
    static a6: f64 = -0.1384794;
    static a7: f64 = 0.1250060;

    static one_7: f64 = 0.1428571428571428571;
    static one_12: f64 = 0.0833333333333333333;
    static one_24: f64 = 0.0416666666666666667;

    /* These are static --- persistent between calls for same mu : */
    let mut l = L.with(|L_| *L_.borrow()); /* integer "w/o overflow" */
    let mut m = M.with(|M_| *M_.borrow()); /*, muold  = 0.*/
    let mut b1 = B1.with(|B1_| *B1_.borrow());
    let mut b2 = B2.with(|B2_| *B2_.borrow());
    let mut c = C.with(|C_| *C_.borrow());
    let mut c0 = C0.with(|C0_| *C0_.borrow());
    let mut c1 = C1.with(|C1_| *C1_.borrow());
    let mut c2 = C2.with(|C2_| *C2_.borrow());
    let mut c3 = C3.with(|C3_| *C3_.borrow());
    let mut pp = PP.with(|PP_| *PP_.borrow());
    let mut p0 = P0.with(|P0_| *P0_.borrow());
    let mut p = P.with(|P_| *P_.borrow());
    let mut q = Q.with(|Q_| *Q_.borrow());
    let mut s = S.with(|S_| *S_.borrow());
    let mut d = D.with(|D_| *D_.borrow());
    let mut omega = OMEGA.with(|OMEGA_| *OMEGA_.borrow());
    let mut big_l = BIG_L.with(|BIG_L_| *BIG_L_.borrow());
    let mut muprev = MUPREV.with(|MUPREV_| *MUPREV_.borrow());
    let mut muprev2 = MUPREV2.with(|MUPREV2_| *MUPREV2_.borrow());

    /* Local Vars  [initialize some for -Wall]: */
    let mut del = 0.0; /* end {initialize persistent vars} */
    let mut difmuk = 0.0;
    let mut E = 0.0;
    let mut fk = 0.0;
    let mut fx = 0.0;
    let mut fy = 0.0;
    let mut g = 0.0;
    let mut px = 0.0;
    let mut py = 0.0;
    let mut t = 0.0;
    let mut u = 0.0;
    let mut v = 0.0;
    let mut x = 0.0;
    let mut pois = -1.0;
    let mut k = 0;
    let mut kflag = 0;
    let mut big_mu = false;
    let mut new_big_mu = false;

    if !mu.is_finite() || mu < 0.0 {
        return f64::nan().into();
    }
    if mu <= 0.0 {
        return 0.0.into();
    }

    big_mu = mu >= 10.0;
    if big_mu {
        new_big_mu = false
    }

    if !(big_mu && mu == muprev) {
        /* maybe compute new persistent par.*s */
        if big_mu {
            new_big_mu = true;
            /* Case A. (recalculation of *s,*d,*l because mu has changed):
             * The poisson probabilities pk exceed the discrete normal
             * probabilities fk whenever k >= *m(mu).
             */
            MUPREV.with(|MUPREV_| {
                *MUPREV_.borrow_mut() = mu;
                muprev = *MUPREV_.borrow();
            });
            S.with(|S_| {
                *S_.borrow_mut() = mu.sqrt();
                s = *S_.borrow();
            });
            D.with(|D_| {
                *D_.borrow_mut() = 6.0 * mu * mu;
                d = *D_.borrow();
            });
            BIG_L.with(|BIG_L_| {
                *BIG_L_.borrow_mut() = (mu - 1.1484).floor();
                big_l = *BIG_L_.borrow();
            });
        /* = an upper bound to *m(mu) for all mu >= 10.*/
        } else {
            /* Small mu ( < 10) -- not using normal approx. */

            /* Case B. (start new table and calculate *p0 if necessary) */

            /* muprev = 0.0;-* such that next time, mu != *muprev ..*/
            if mu != muprev {
                MUPREV.with(|MUPREV_| {
                    *MUPREV_.borrow_mut() = mu;
                    muprev = *MUPREV_.borrow();
                }); /* *pp[] is already ok up to *pp[*l] */
                M.with(|M_| {
                    *M_.borrow_mut() = 1.max(mu as i32);
                    m = *M_.borrow();
                });
                L.with(|L_| {
                    *L_.borrow_mut() = 0;
                    l = *L_.borrow();
                });
                P.with(|P_| {
                    *P_.borrow_mut() = (-mu).exp();
                    p = *P_.borrow();
                });
                P0.with(|P0_| {
                    *P0_.borrow_mut() = p;
                    p0 = *P0_.borrow();
                });
                Q.with(|Q_| {
                    *Q_.borrow_mut() = p0;
                    q = *Q_.borrow();
                });
            }
            loop {
                /* Step U. uniform sample for inversion method */
                u = rng.unif_rand();
                if u <= p0 {
                    return 0.0.into();
                }
                /* Step T. table comparison until the end *pp[*l] of the
                *pp-table of cumulative poisson probabilities
                (0.458 > ~= *pp[9](= 0.45792971447) for mu=10 ) */
                if l != 0 {
                    k = if u <= 0.458 { 1 } else { l.min(m as usize) };
                    while k <= l {
                        if u <= pp[k] {
                            return k.into();
                        }
                        k += 1
                    }
                    if l == 35 {
                        /* u > *pp[35] */
                        continue;
                    }
                }
                /* Step C. creation of new poisson
                probabilities *p[*l..] and their cumulatives *q =: *pp[k] */
                L.with(|L_| {
                    *L_.borrow_mut() += 1;
                    l = *L_.borrow();
                });
                k = l;
                while k <= 35 {
                    P.with(|P_| {
                        *P_.borrow_mut() *= mu / k as f64;
                        p = *P_.borrow();
                    });
                    Q.with(|Q_| {
                        *Q_.borrow_mut() += p;
                        q = *Q_.borrow();
                    });
                    PP.with(|PP_| {
                        PP_.borrow_mut()[k] = q;
                        pp[k] = PP_.borrow()[k];
                    });
                    if u <= q {
                        L.with(|L_| {
                            *L_.borrow_mut() = k;
                            l = *L_.borrow();
                        });
                        return k.into();
                    }
                    k += 1
                }
                L.with(|L_| {
                    *L_.borrow_mut() = 35;
                    l = *L_.borrow();
                });
            }
        }
        /* mu < 10 */
    }

    /* Only if mu >= 10 : ----------------------- */

    /* Step N. normal sample */
    g = mu + s * rng.norm_rand(); /* norm_rand() ~ N(0,1), standard normal */
    if g >= 0.0 {
        pois = g.floor();
        /* Step I. immediate acceptance if pois is large enough */
        if pois >= big_l {
            return pois.into();
        }
        /* Step S. squeeze acceptance */
        fk = pois; /* ~ U(0,1) - sample */
        difmuk = mu - fk;
        u = rng.unif_rand();
        if d * u >= difmuk * difmuk * difmuk {
            return pois.into();
        }
    }

    /* Step P. preparations for steps Q and H.
    (recalculations of parameters if necessary) */

    if new_big_mu || mu != muprev2 {
        /* Careful! *muprev2 is not always == *muprev
        because one might have exited in step I or S
        */
        MUPREV2.with(|MUPREV2_| {
            *MUPREV2_.borrow_mut() = mu;
            muprev2 = *MUPREV2_.borrow();
        });
        OMEGA.with(|OMEGA_| {
            *OMEGA_.borrow_mut() = strafe_consts::_1DSQRT_2TPI / s;
            omega = *OMEGA_.borrow();
        });
        /* The quantities *b1, *b2, *c3, *c2, *c1, *c0 are for the Hermite
         * approximations to the discrete normal probabilities fk. */
        /* 'Subroutine' F is called (kflag=0 for correct return) */

        B1.with(|B1_| {
            *B1_.borrow_mut() = one_24 / mu;
            b1 = *B1_.borrow();
        });
        B2.with(|B2_| {
            *B2_.borrow_mut() = 0.3 * b1 * b1;
            b2 = *B2_.borrow();
        });
        C3.with(|C3_| {
            *C3_.borrow_mut() = one_7 * b1 * b2;
            c3 = *C3_.borrow();
        });
        C2.with(|C2_| {
            *C2_.borrow_mut() = b2 - 15.0 * c3;
            c2 = *C2_.borrow();
        });
        C1.with(|C1_| {
            *C1_.borrow_mut() = b1 - 6.0 * b2 + 45.0 * c3;
            c1 = *C1_.borrow();
        });
        C0.with(|C0_| {
            *C0_.borrow_mut() = 1.0 - b1 + 3.0 * b2 - 15.0 * c3;
            c0 = *C0_.borrow();
        });
        /* guarantees majorization by the 'hat'-function. */
        C.with(|C_| {
            *C_.borrow_mut() = 0.1069 / mu;
            c = *C_.borrow();
        });
    }

    if g >= 0.0 {
        kflag = 0;
        current_block = 14971143117904085953;
    } else {
        current_block = 5722677567366458307;
    }

    loop {
        match current_block {
            14971143117904085953 => {
                /* 'subroutine' F : calculation of px,py,fx,fy. */
                if pois < 10.0 {
                    /* use factorials from table fact[] */
                    px = -mu;
                    py = mu.powf(pois) / fact[pois as usize]
                } else {
                    /* Case pois >= 10 uses polynomial approximation
                    a0-a7 for accuracy when advisable */
                    del = one_12 / fk;
                    del = del * (1.0 - 4.8 * del * del);
                    v = difmuk / fk;
                    if v.abs() <= 0.25 {
                        px = fk
                            * v
                            * v
                            * (((((((a7 * v + a6) * v + a5) * v + a4) * v + a3) * v + a2) * v + a1)
                                * v
                                + a0)
                            - del
                    } else {
                        /* |v| > 1/4 */
                        px = fk * (1.0 + v).ln() - difmuk - del
                    } /* x^2 */
                    py = strafe_consts::_1DSQRT_2TPI / fk.sqrt()
                }
                x = (0.5 - difmuk) / s;
                x *= x;
                fx = -0.5 * x;
                fy = omega * (((c3 * x + c2) * x + c1) * x + c0);
                if kflag > 0 {
                    /* Step H. Hat acceptance (E is repeated on rejection) */
                    if c * u.abs() <= py * (px + E).exp() - fy * (fx + E).exp() {
                        break;
                    } else {
                        current_block = 5722677567366458307;
                    }
                } else if fy - u * fy <= py * (px - fx).exp()
                /* Step Q. Quotient acceptance (rare case) */
                {
                    break;
                } else {
                    current_block = 5722677567366458307;
                }
            }
            _ => {
                /* Step E. Exponential Sample */
                E = rng.exp_rand(); /* ~ Exp(1) (standard exponential) */
                /* t > -.67.. */
                /*  sample t from the laplace 'hat'
                (if t <= -0.6744 then pk < fk for all mu >= 10.0) */
                u = 2.0 * rng.unif_rand() - 1.0;
                t = 1.8 + (E.abs() * u.signum());
                if !(t > -0.6744) {
                    current_block = 5722677567366458307;
                    continue;
                }
                pois = (mu + s * t).floor();
                fk = pois;
                difmuk = mu - fk;
                /* 'subroutine' F is called (kflag=1 for correct return) */
                kflag = 1;
                current_block = 14971143117904085953;
            }
        }
    }

    pois.into()
}