r2rs-nmath 0.1.1

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

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

/// Beta density,
///
/// $ p(x;a,b) = \frac{(a+b-1)!}{(a-1)!(b-1)!}x^{a-1}(1-x)^{b-1} $
///
/// The basic formula for the log density is thus
/// $(a-1) log x + (b-1) log (1-x) - lbeta(a, b)$
/// If either ```a``` or ```b <= 2``` then ```0 < lbeta(a, b) < 710``` and so no
/// term is large.  We use Loader's code only if both ```a``` and ```b > 2```.
pub fn dbeta<R: Into<Real64>, P1: Into<Positive64>, P2: Into<Positive64>>(
    x: R,
    a: P1,
    b: P2,
    log: bool,
) -> Real64 {
    let x = x.into().unwrap();
    let a = a.into().unwrap();
    let b = b.into().unwrap();

    if !(0.0..=1.0).contains(&x) {
        return f64::d_0(log).into();
    }

    // limit cases for (a,b), leading to point masses
    if a == 0.0 || b == 0.0 || !a.is_finite() || !b.is_finite() {
        if a == 0.0 && b == 0.0 {
            // point mass 1/2 at each of {0,1} :
            return if x == 0.0 || x == 1.0 {
                f64::infinity()
            } else {
                f64::d_0(log)
            }
            .into();
        }
        if a == 0.0 || a / b == 0.0 {
            // point mass 1 at 0
            return if x == 0.0 {
                f64::infinity()
            } else {
                f64::d_0(log)
            }
            .into();
        }
        if b == 0.0 || b / a == 0.0 {
            // point mass 1 at 1
            return if x == 1.0 {
                f64::infinity()
            } else {
                f64::d_0(log)
            }
            .into();
        }
        // else, remaining case:  a = b = Inf : point mass 1 at 1/2
        return if x == 0.5 {
            f64::infinity()
        } else {
            f64::d_0(log)
        }
        .into();
    }

    if x == 0.0 {
        if a > 1.0 {
            return f64::d_0(log).into();
        }
        if a < 1.0 {
            return f64::infinity().into();
        }
        /* a == 1 : */
        return b.d_val(log).into();
    }
    if x == 1.0 {
        if b > 1.0 {
            return f64::d_0(log).into();
        }
        if b < 1.0 {
            return f64::infinity().into();
        }
        /* b == 1 : */
        return a.d_val(log).into();
    }

    let mut lval = 0.0;
    if a <= 2.0 || b <= 2.0 {
        lval = (a - 1.0) * x.ln() + (b - 1.0) * (-x).ln_1p() - lbeta(a, b).unwrap()
    } else {
        lval = (a + b - 1.0).ln() + dbinom_raw(a - 1.0, a + b - 2.0, x, 1.0 - x, true)
    }

    return lval.d_exp(log).into();
}