use num_traits::Float;
use strafe_type::{FloatConstraint, LogProbability64, PositiveInteger64, Probability64, Real64};
use super::{log_pbinom, pbinom};
use crate::{distribution::func::discrete_body, traits::DPQ};
pub fn qbinom<PR1: Into<Probability64>, PO: Into<PositiveInteger64>, PR2: Into<Probability64>>(
p: PR1,
n: PO,
pr: PR2,
lower_tail: bool,
) -> Real64 {
let p = p.into().unwrap();
qbinom_inner(p, n, pr, lower_tail, false)
}
pub fn log_qbinom<
LP: Into<LogProbability64>,
PO: Into<PositiveInteger64>,
PR: Into<Probability64>,
>(
p: LP,
n: PO,
pr: PR,
lower_tail: bool,
) -> Real64 {
let p = p.into().unwrap();
qbinom_inner(p, n, pr, lower_tail, true)
}
fn qbinom_inner<PO: Into<PositiveInteger64>, PR: Into<Probability64>>(
p: f64,
n: PO,
pr: PR,
lower_tail: bool,
log: bool,
) -> Real64 {
let n = n.into().unwrap();
let pr = pr.into().unwrap();
if !n.is_finite() || !pr.is_finite() {
return f64::nan().into();
}
if !p.is_finite() && !log {
return f64::nan().into();
}
if let Some(ret) = p.q_p01_boundaries(0.0, n, lower_tail, log) {
return ret.into();
}
if pr == 0.0 || n == 0.0 {
return 0.0.into();
}
if pr == 1.0 {
return n.into();
}
let q = 1.0 - pr;
let mu = n * pr;
let sigma = (n * pr * q).sqrt();
let gamma = (q - pr) / sigma;
let ret = discrete_body(
mu,
sigma,
gamma,
p,
Some(n),
lower_tail,
log,
&|p| pbinom(p, n, pr, lower_tail).unwrap(),
&|p| log_pbinom(p, n, pr, lower_tail).unwrap(),
);
ret.into()
}