//! Provides functions related to factorial calculations (e.g. binomial
//! coefficient, factorial, multinomial)

use crate::error::StatsError;
use crate::function::gamma;
use crate::Result;
use core::f64::INFINITY as INF;

/// The maximum factorial representable
/// by a 64-bit floating point without
/// overflowing
pub const MAX_FACTORIAL: usize = 170;

/// Computes the factorial function `x -> x!` for
/// `170 >= x >= 0`. All factorials larger than `170!`
/// will overflow an `f64`.
///
/// # Remarks
///
/// Returns `f64::INFINITY` if `x > 170`
pub fn factorial(x: u64) -> f64 {
    let x = x as usize;
    FCACHE.get(x).map_or(INF, |&fac| fac)
}

/// Computes the logarithmic factorial function `x -> ln(x!)`
/// for `x >= 0`.
///
/// # Remarks
///
/// Returns `0.0` if `x <= 1`
pub fn ln_factorial(x: u64) -> f64 {
    let x = x as usize;
    FCACHE
        .get(x)
        .map_or_else(|| gamma::ln_gamma(x as f64 + 1.0), |&fac| fac.ln())
}

/// Computes the binomial coefficient `n choose k`
/// where `k` and `n` are non-negative values.
///
/// # Remarks
///
/// Returns `0.0` if `k > n`
pub fn binomial(n: u64, k: u64) -> f64 {
    if k > n {
        0.0
    } else {
        (0.5 + (ln_factorial(n) - ln_factorial(k) - ln_factorial(n - k)).exp()).floor()
    }
}

/// Computes the natural logarithm of the binomial coefficient
/// `ln(n choose k)` where `k` and `n` are non-negative values
///
/// # Remarks
///
/// Returns `f64::NEG_INFINITY` if `k > n`
pub fn ln_binomial(n: u64, k: u64) -> f64 {
    if k > n {
        f64::NEG_INFINITY
    } else {
        ln_factorial(n) - ln_factorial(k) - ln_factorial(n - k)
    }
}

/// Computes the multinomial coefficient: `n choose n1, n2, n3, ...`
///
/// # Panics
///
/// If the elements in `ni` do not sum to `n`
pub fn multinomial(n: u64, ni: &[u64]) -> f64 {
    checked_multinomial(n, ni).unwrap()
}

/// Computes the multinomial coefficient: `n choose n1, n2, n3, ...`
///
/// # Errors
///
/// If the elements in `ni` do not sum to `n`
pub fn checked_multinomial(n: u64, ni: &[u64]) -> Result<f64> {
    let (sum, ret) = ni.iter().fold((0, ln_factorial(n)), |acc, &x| {
        (acc.0 + x, acc.1 - ln_factorial(x))
    });
    if sum != n {
        Err(StatsError::ContainerExpectedSumVar("ni", "n"))
    } else {
        Ok((0.5 + ret.exp()).floor())
    }
}

// Initialization for pre-computed cache of 171 factorial
// values 0!...170!
lazy_static! {
    static ref FCACHE: [f64; MAX_FACTORIAL + 1] = {
        let mut fcache = [1.0; MAX_FACTORIAL + 1];
        fcache
            .iter_mut()
            .enumerate()
            .skip(1)
            .fold(1.0, |acc, (i, elt)| {
                let fac = acc * i as f64;
                *elt = fac;
                fac
            });
        fcache
    };
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_factorial_and_ln_factorial() {
        let mut fac = 1.0;
        assert_eq!(factorial(0), fac);
        for i in 1..171 {
            fac *= i as f64;
            assert_eq!(factorial(i), fac);
            assert_eq!(ln_factorial(i), fac.ln());
        }
    }

    #[test]
    fn test_factorial_overflow() {
        assert_eq!(factorial(172), INF);
        assert_eq!(factorial(u64::MAX), INF);
    }

    #[test]
    fn test_ln_factorial_does_not_overflow() {
        assert_eq!(ln_factorial(1 << 10), 6078.2118847500501140);
        assert_almost_eq!(ln_factorial(1 << 12), 29978.648060844048236, 1e-11);
        assert_eq!(ln_factorial(1 << 15), 307933.81973375485425);
        assert_eq!(ln_factorial(1 << 17), 1413421.9939462073242);
    }

    #[test]
    fn test_binomial() {
        assert_eq!(binomial(1, 1), 1.0);
        assert_eq!(binomial(5, 2), 10.0);
        assert_eq!(binomial(7, 3), 35.0);
        assert_eq!(binomial(1, 0), 1.0);
        assert_eq!(binomial(0, 1), 0.0);
        assert_eq!(binomial(5, 7), 0.0);
    }

    #[test]
    fn test_ln_binomial() {
        assert_eq!(ln_binomial(1, 1), 1f64.ln());
        assert_almost_eq!(ln_binomial(5, 2), 10f64.ln(), 1e-14);
        assert_almost_eq!(ln_binomial(7, 3), 35f64.ln(), 1e-14);
        assert_eq!(ln_binomial(1, 0), 1f64.ln());
        assert_eq!(ln_binomial(0, 1), 0f64.ln());
        assert_eq!(ln_binomial(5, 7), 0f64.ln());
    }

    #[test]
    fn test_multinomial() {
        assert_eq!(1.0, multinomial(1, &[1, 0]));
        assert_eq!(10.0, multinomial(5, &[3, 2]));
        assert_eq!(10.0, multinomial(5, &[2, 3]));
        assert_eq!(35.0, multinomial(7, &[3, 4]));
    }

    #[test]
    #[should_panic]
    fn test_multinomial_bad_ni() {
        multinomial(1, &[1, 1]);
    }

    #[test]
    fn test_checked_multinomial_bad_ni() {
        assert!(checked_multinomial(1, &[1, 1]).is_err());
    }
}