r2rs-nmath 0.1.1

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

use crate::{distribution::binom::dbinom_raw, traits::DPQ};

/// Computes the geometric probabilities, Pr(X=x) = p(1-p)^x.
pub fn dgeom<R: Into<Real64>, P: Into<Probability64>>(x: R, p: P, log: bool) -> Real64 {
    let mut x = x.into().unwrap();
    let p = p.into().unwrap();

    let mut prob = 0.0;

    if p <= 0.0 {
        return f64::nan().into();
    }

    if x.is_non_integer() {
        warn!("non-integer x = {}", x);
        return f64::d_0(log).into();
    }

    if x < 0.0 || !x.is_finite() || p == 0.0 {
        return f64::d_0(log).into();
    }
    x = x.round();

    /* prob = (1-p)^x, stable for small p */
    prob = dbinom_raw(0.0, x, p, 1.0 - p, log);
    if log { (p.ln()) + prob } else { (p) * prob }.into()
}