r2rs-nmath 0.1.1

// Translation of nmath's rbeta
//
// 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 strafe_type::{FloatConstraint, Positive64, Real64};

use crate::traits::RNG;

thread_local! {
    static BETA: RefCell<f64> = RefCell::new(0.0);
    static GAMMA: RefCell<f64> = RefCell::new(0.0);
    static DELTA: RefCell<f64> = RefCell::new(0.0);
    static K1: RefCell<f64> = RefCell::new(0.0);
    static K2: RefCell<f64> = RefCell::new(0.0);
    static OLDA: RefCell<f64> = RefCell::new(-1.0);
    static OLDB: RefCell<f64> = RefCell::new(-1.0);
}

pub fn reset_rbeta_statics() {
    BETA.with(|BETA_| *BETA_.borrow_mut() = 0.0);
    GAMMA.with(|GAMMA_| *GAMMA_.borrow_mut() = 0.0);
    DELTA.with(|DELTA_| *DELTA_.borrow_mut() = 0.0);
    K1.with(|K1_| *K1_.borrow_mut() = 0.0);
    K2.with(|K2_| *K2_.borrow_mut() = 0.0);
    OLDA.with(|OLDA_| *OLDA_.borrow_mut() = -1.0);
    OLDB.with(|OLDB_| *OLDB_.borrow_mut() = -1.0);
}

const EXPMAX: f64 = f64::MAX_EXP as f64 * std::f64::consts::LN_2;

/// Random beta variates
pub fn rbeta<P1, P2, R>(aa: P1, bb: P2, rng: &mut R) -> Real64
where
    P1: Into<Positive64>,
    P2: Into<Positive64>,
    R: RNG,
{
    let aa = aa.into().unwrap();
    let bb = bb.into().unwrap();

    if !aa.is_finite() && !bb.is_finite() {
        // a = b = Inf : all mass at 1/2
        return 0.5.into();
    }
    if aa == 0.0 && bb == 0.0 {
        // point mass 1/2 at each of {0,1} :
        return if rng.unif_rand() < 0.5 { 0.0 } else { 1.0 }.into();
    }
    // now, at least one of a, b is finite and positive
    if !aa.is_finite() || bb == 0.0 {
        return 1.0.into();
    }
    if !bb.is_finite() || aa == 0.0 {
        return 0.0.into();
    }

    let mut a = 0.0;
    let mut b = 0.0;
    let mut alpha = 0.0;
    let mut r = 0.0;
    let mut s = 0.0;
    let mut t = 0.0;
    let mut u1 = 0.0;
    let mut u2 = 0.0;
    let mut v = 0.0;
    let mut w = 0.0;
    let mut y = 0.0;
    let mut z = 0.0;
    /* Uses these GLOBALS to save time when many rv's are generated : */
    let mut beta = BETA.with(|BETA_| *BETA_.borrow());
    let mut gamma = GAMMA.with(|GAMMA_| *GAMMA_.borrow());
    let mut delta = DELTA.with(|DELTA_| *DELTA_.borrow());
    let mut k1 = K1.with(|K1_| *K1_.borrow());
    let mut k2 = K2.with(|K2_| *K2_.borrow());
    let mut olda = OLDA.with(|OLDA_| *OLDA_.borrow());
    let mut oldb = OLDB.with(|OLDB_| *OLDB_.borrow());

    /* Test if we need new "initializing" */
    let qsame = olda == aa && oldb == bb; /* a <= b */
    if !qsame {
        OLDA.with(|OLDA_| {
            *OLDA_.borrow_mut() = aa;
            olda = *OLDA_.borrow();
        });
        OLDB.with(|OLDB_| {
            *OLDB_.borrow_mut() = bb;
            oldb = *OLDB_.borrow();
        });
    }

    a = aa.min(bb);
    b = aa.max(bb);
    alpha = a + b;

    let v_w_from__u1_bet = |AA: f64, beta: &mut f64, u1: &mut f64, w: &mut f64, v: &mut f64| {
        *v = *beta * (*u1 / (1.0_f64 - *u1)).ln();
        if *v <= EXPMAX {
            *w = AA * v.exp();
            if !w.is_finite() {
                *w = f64::MAX
            }
        } else {
            *w = f64::MAX
        }
    };

    let ret = if a <= 1.0 {
        /* --- Algorithm BC --- */
        /* changed notation, now also a <= b (was reversed) */
        if !qsame {
            /* initialize */
            BETA.with(|BETA_| {
                *BETA_.borrow_mut() = 1.0 / a;
                beta = *BETA_.borrow();
            });
            DELTA.with(|DELTA_| {
                *DELTA_.borrow_mut() = 1.0 + b - a;
                delta = *DELTA_.borrow();
            });
            K1.with(|K1_| {
                *K1_.borrow_mut() = delta * (0.0138889 + 0.0416667 * a) / (b * beta - 0.777778);
                k1 = *K1_.borrow();
            });
            K2.with(|K2_| {
                *K2_.borrow_mut() = 0.25 + (0.5 + 0.25 / delta) * a;
                k2 = *K2_.borrow();
            });
        }
        loop
        /* FIXME: "do { } while()", but not trivially because of "continue"s:*/
        {
            u1 = rng.unif_rand();
            u2 = rng.unif_rand();
            if u1 < 0.5 {
                y = u1 * u2;
                z = u1 * y;
                if 0.25 * u2 + z - y >= k1 {
                    continue;
                }
            } else {
                z = u1 * u1 * u2;
                if z <= 0.25 {
                    v_w_from__u1_bet(b, &mut beta, &mut u1, &mut w, &mut v);
                    break;
                }
                if z >= k2 {
                    continue;
                }
            }

            v_w_from__u1_bet(b, &mut beta, &mut u1, &mut w, &mut v);

            if alpha * ((alpha / (a + w)).ln() + v) - 1.3862944 >= z.ln() {
                break;
            }
        }
        if aa == a {
            a / (a + w)
        } else {
            w / (a + w)
        }
    } else {
        /* Algorithm BB */

        if !qsame {
            /* initialize */
            BETA.with(|BETA_| {
                *BETA_.borrow_mut() = ((alpha - 2.0) / (2.0 * a * b - alpha)).sqrt();
                beta = *BETA_.borrow();
            });
            GAMMA.with(|GAMMA_| {
                *GAMMA_.borrow_mut() = a + 1.0 / beta;
                gamma = *GAMMA_.borrow();
            });
        }
        loop {
            u1 = rng.unif_rand();
            u2 = rng.unif_rand();

            v_w_from__u1_bet(a, &mut beta, &mut u1, &mut w, &mut v);

            z = u1 * u1 * u2;
            r = gamma * v - 1.3862944;
            s = a + r - w;
            if s + 2.609438 >= 5.0 * z {
                break;
            }
            t = z.ln();
            if s > t {
                break;
            }
            if !(r + alpha * (alpha / (b + w)).ln() < t) {
                break;
            }
        }

        if aa != a {
            b / (b + w)
        } else {
            w / (b + w)
        }
    };
    return ret.into();
}