r2rs-nmath 0.1.1

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

use crate::{func::lfastchoose, traits::DPQ};

/// The quantile function of the hypergeometric distribution.
pub fn qhyper<
    PR: Into<Probability64>,
    PI1: Into<PositiveInteger64>,
    PI2: Into<PositiveInteger64>,
    PI3: Into<PositiveInteger64>,
>(
    p: PR,
    NR: PI1,
    NB: PI2,
    n: PI3,
    lower_tail: bool,
) -> Real64 {
    let p = p.into().unwrap();
    qhyper_inner(p, NR, NB, n, lower_tail, false)
}

/// The quantile function of the hypergeometric distribution.
pub fn log_qhyper<
    LP: Into<LogProbability64>,
    PI1: Into<PositiveInteger64>,
    PI2: Into<PositiveInteger64>,
    PI3: Into<PositiveInteger64>,
>(
    p: LP,
    NR: PI1,
    NB: PI2,
    n: PI3,
    lower_tail: bool,
) -> Real64 {
    let p = p.into().unwrap();
    qhyper_inner(p, NR, NB, n, lower_tail, true)
}

fn qhyper_inner<
    P1: Into<PositiveInteger64>,
    P2: Into<PositiveInteger64>,
    P3: Into<PositiveInteger64>,
>(
    mut p: f64,
    NR: P1,
    NB: P2,
    n: P3,
    lower_tail: bool,
    log: bool,
) -> Real64 {
    let mut NR = NR.into().unwrap();
    let mut NB = NB.into().unwrap();
    let mut n = n.into().unwrap();

    /* This is basically the same code as  ./phyper.c  *used* to be --> FIXME! */
    let mut N = 0.0;
    let mut xstart = 0.0;
    let mut xend = 0.0;
    let mut xr = 0.0;
    let mut xb = 0.0;
    let mut sum = 0.0;
    let mut term = 0.0;
    let mut small_N = false;

    if !p.is_finite() || !NR.is_finite() || !NB.is_finite() || !n.is_finite() {
        return f64::nan().into();
    }

    NR = NR.round();
    NB = NB.round();
    N = NR + NB;
    n = n.round();
    if n > N {
        return f64::nan().into();
    }

    /* Goal:  Find  xr (= #{red balls in sample}) such that
     *   phyper(xr,  NR,NB, n) >= p > phyper(xr - 1,  NR,NB, n)
     */

    xstart = (n - NB).max(0.0); /* always ( = #{black balls in sample} ) */
    xend = n.min(NR); /* won't have underflow in product below */

    if let Some(ret) = p.q_p01_boundaries(xstart, xend, lower_tail, log) {
        return ret.into();
    }

    xr = xstart;
    xb = n - xr;

    small_N = N < 1000.0;
    /* if N is small,  term := product.ratio( bin.coef );
    otherwise work with its logarithm to protect against underflow */
    if xr > NR + 1.0 || xb > NB + 1.0 || n > N + 1.0 {
        term = f64::nan();
    } else {
        term = lfastchoose(NR, xr).unwrap() + lfastchoose(NB, xb).unwrap()
            - lfastchoose(N, n).unwrap(); /* was 64, but failed on FreeBSD sometimes */
    }
    if small_N {
        term = term.exp()
    }
    NR -= xr;
    NB -= xb;

    if log || !lower_tail {
        p = p.dt_qiv(lower_tail, log);
    }
    p *= 1.0 - 1000.0 * 2.2204460492503131e-16;
    sum = if small_N { term } else { term.exp() };

    while sum < p && xr < xend {
        xr += 1.0;
        NB += 1.0;
        if small_N {
            term *= (NR / xr) * (xb / NB)
        } else {
            term += ((NR / xr) * (xb / NB)).ln()
        }
        sum += if small_N { term } else { term.exp() };
        xb -= 1.0;
        NR -= 1.
    }
    xr.into()
}