r2rs-nmath 0.1.1

// Translation of nmath's qbinom
//
// 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 num_traits::Float;
use strafe_type::FloatConstraint;

use crate::distribution::norm::{log_qnorm, qnorm};

fn do_search(
    mut y: f64,
    z: &mut f64,
    p: f64,
    dist_max_y: Option<f64>,
    incr: f64,
    lower_tail: bool,
    log: bool,
    p_func: &dyn Fn(f64) -> f64,
    log_p_func: &dyn Fn(f64) -> f64,
) -> f64 {
    {
        let left = if lower_tail { *z >= p } else { *z < p };
        // R_DBG_printf(" do_search(y=%g, z=%15.10g %s p = %15.10g) --> search to %s",
        //              y, *z, (lower_tail ? ">=" : "<"), p, (left ? "left" : "right"));
        // if incr > 1.0 {
        //     R_DBG_printf(", incr = %.0f\n", incr);
        // }
        // else {
        //     R_DBG_printf("\n");
        // }

        if left {
            // (lower_tail, *z >= p)  or  (upper tail, *z < p): search to the __left__
            // let mut iter = 0;
            loop {
                let mut newz = -1.0; // -Wall
                if y > 0.0 {
                    newz = if log {
                        log_p_func(y - incr)
                    } else {
                        p_func(y - incr)
                    };
                } else if y < 0.0 {
                    y = 0.0;
                }
                // note that newz may be NaN because of remaining border line bugs in _pDIST_() {bug from pbeta()}
                if y == 0.0 || newz.is_nan() || if lower_tail { newz < p } else { newz >= p } {
                    // R_DBG_printf("  new y=%.15g, " AS_CHAR(_pDIST_) "(y-incr,*) %s;"
                    //              " ==> search() returns previous z=%g after %d iter.\n", y,
                    //              ISNAN(newz) ? "is NaN" : (lower_tail ? "< p" : ">= p"), *z, iter);
                    return y; // and previous *z
                }
                y = (y - incr).max(0.0);
                *z = newz;
                // iter += 1;
            }
        } else {
            // (lower_tail, *z < p)  or  (upper tail, *z >= p): search to the __right__
            // let mut iter = 0;
            loop {
                y += incr;
                if let Some(y_max) = dist_max_y {
                    if y < y_max {
                        *z = if log { log_p_func(y) } else { p_func(y) };
                    } else if y > y_max {
                        y = y_max
                    }
                } else {
                    *z = if log { log_p_func(y) } else { p_func(y) };
                }

                let dist_bool = if let Some(y_max) = dist_max_y {
                    y == y_max
                } else {
                    false
                };

                if dist_bool || z.is_nan() || if lower_tail { *z >= p } else { *z < p } {
                    // R_DBG_printf("  new y=%.15g, z=%g = " AS_CHAR(_pDIST_) "(y,*) %s;"
                    //              " ==> search() returns after %d iter.\n", y, *z,
                    //              ISNAN(*z) ? "is NaN" : (lower_tail ? ">= p" : "< p"), iter);
                    return y;
                }
                // iter += 1;
            }
        }
    }
}

fn discr_check_boundary(y: &mut f64, dist_max_y: Option<f64>) {
    if let Some(max_y) = dist_max_y {
        if *y > max_y {
            // way off
            *y = max_y;
        }
    }
    if *y < 0.0 {
        *y = 0.0;
    }
}

pub fn discrete_body(
    mu: f64,
    sigma: f64,
    gamma: f64,
    mut p: f64,
    dist_max_y: Option<f64>,
    lower_tail: bool,
    log: bool,
    p_func: &dyn Fn(f64) -> f64,
    log_p_func: &dyn Fn(f64) -> f64,
) -> f64 {
    let mut incr = 0.0;
    /* y := approx.value (Cornish-Fisher expansion) :  */
    let mut z = if log {
        log_qnorm(p, 0.0, 1.0, lower_tail).unwrap()
    } else {
        qnorm(p, 0.0, 1.0, lower_tail).unwrap()
    };
    // TODO: different than R's nearbyint function usage for 0.5, but Rust correctly gets 1.0 whereas C gets 0.0
    let mut y = (mu + sigma * (z + gamma * (z * z - 1.0) / 6.0)).round();
    // R_DBG_printf(" Cornish-Fisher: initial z = qnorm(p, l.t, log)= %g,  y = %g;\n", z,y);

    discr_check_boundary(&mut y, dist_max_y);

    z = if log { log_p_func(y) } else { p_func(y) };

    /* Algorithmic "tuning parameters", used to be hardwired; changed for speed &| precision */
    let _pf_n_ = 8; /* was hardwired to 64 */
    let _pf_L_ = 2; /* was hardwired to 64 */
    let _yLarge_ = 4096; /* was hardwired to 1e5 */
    let _incF_ = 1.0 / 64.0; /* was hardwired to 0.001 (= 1./1000 ) */
    let _iShrink_ = 8; /* was hardwired to  100 */
    let _relTol_ = 1e-15; /* was hardwired to 1e-15 */
    let _xf_ = 4; /* extra factor, *must* be >= 1 (new anyway) */

    // R_DBG_printf(" algo. tuning: fuzz factors _pf_{n,L}: {%.0f, %.0f};  yLarge = %g\n"
    // "      large case: _incF_=%g, _iShrink_=%g; _relTol_=%g, _xf_=%g\n",
    // _pf_n_, _pf_L_, _yLarge_,   _incF_, _iShrink_, _relTol_, _xf_);

    /* fuzz to ensure left continuity: do not loose too much (=> error in upper tail) */
    if log {
        /* <==> p \in [-Inf, 0]  different adjustment: "other sign" */
        let e = _pf_L_ as f64 * f64::EPSILON;
        if lower_tail && p > -f64::max_value() {
            /* prevent underflow to -Inf */
            p *= 1.0 + e;
        } else {
            /* if(p < - DBL_MIN) // not too close to -0 */
            p *= 1.0 - e;
        }
    } else {
        /* not log scale */
        let e = _pf_n_ as f64 * f64::EPSILON;
        if lower_tail {
            p *= 1.0 - e;
        } else if 1.0 - p > _xf_ as f64 * e {
            /* otherwise get p > 1 */
            p *= 1.0 + e;
        }
    }
    // R_DBG_printf("  new z := " AS_CHAR(_pDIST_) "(y, *) = %.11g; left-cont. fuzz => p = %.11g\n", z, p);

    /* If the C-F value  y   is not too large a simple search is OK */
    if y < _yLarge_ as f64 {
        do_search(
            y, &mut 1.0, p, dist_max_y, incr, lower_tail, log, p_func, log_p_func,
        )
    }
    /* Otherwise be a bit cleverer in the search: use larger increments, notably initially: */
    else {
        /* y >= _yLarge_ */
        let mut oldincr;
        incr = (y * _incF_).floor();
        // let qIt = 0;

        // R_DBG_printf(" large y: --> use larger increments than 1: incr=%.0f\n", incr);
        loop {
            oldincr = incr;
            y = do_search(
                y, &mut z, p, dist_max_y, incr, lower_tail, log, p_func, log_p_func,
            ); /* also updating *z */
            // if(++qIt % 10000 == 0) MAYBE_R_CheckUserInterrupt();
            incr = (incr / _iShrink_ as f64).floor().max(1.0);
            if !(oldincr > 1.0 && incr > y * _relTol_) {
                break;
            }
        }
        // R_DBG_printf("  \\--> oldincr=%.0f, after %d \"outer\" search() iterations\n",
        // oldincr, qIt);
        y
    }
}