numeris 0.5.8

use crate::FloatScalar;
use crate::special::{lgamma, betainc};
use super::{DiscreteDistribution, StatsError};

/// Binomial distribution B(n, p).
///
/// P(X = k) = C(n,k) p^k (1−p)^{n−k} for k = 0, …, n.
///
/// # Example
///
/// ```
/// use numeris::stats::{Binomial, DiscreteDistribution};
///
/// let b = Binomial::new(10, 0.5_f64).unwrap();
/// assert!((b.mean() - 5.0).abs() < 1e-14);
/// assert!((b.variance() - 2.5).abs() < 1e-14);
/// ```
#[derive(Debug, Clone, Copy)]
pub struct Binomial<T> {
    n: u64,
    p: T,
}

impl<T: FloatScalar> Binomial<T> {
    /// Create a binomial distribution with `n` trials and success probability `p`.
    /// Requires `0 ≤ p ≤ 1`.
    pub fn new(n: u64, p: T) -> Result<Self, StatsError> {
        if p < T::zero() || p > T::one() {
            return Err(StatsError::InvalidParameter);
        }
        Ok(Self { n, p })
    }
}

impl<T: FloatScalar> Binomial<T> {
    /// Draw a random sample from this distribution.
    ///
    /// For small n, sums n Bernoulli trials. For large n (> 20 and
    /// np > 5 and n(1-p) > 5), uses a normal approximation with
    /// rounding and clamping.
    pub fn sample(&self, rng: &mut super::Rng) -> u64 {
        let np = T::from(self.n).unwrap() * self.p;
        let nq = T::from(self.n).unwrap() * (T::one() - self.p);
        let five = T::from(5.0).unwrap();

        if self.n <= 20 || np < five || nq < five {
            // Direct summation of Bernoulli trials
            let mut count = 0u64;
            for _ in 0..self.n {
                if rng.next_float::<T>() < self.p {
                    count += 1;
                }
            }
            count
        } else {
            // Normal approximation
            let mean = np;
            let std = (np * (T::one() - self.p)).sqrt();
            let x = mean + std * rng.next_normal::<T>();
            let k = x.round().to_u64().unwrap_or(0);
            k.min(self.n)
        }
    }

    /// Fill a fixed-size array with independent samples.
    pub fn sample_array<const K: usize>(&self, rng: &mut super::Rng) -> [u64; K] {
        let mut out = [0u64; K];
        for v in out.iter_mut() {
            *v = self.sample(rng);
        }
        out
    }
}

impl<T: FloatScalar> DiscreteDistribution<T> for Binomial<T> {
    fn pmf(&self, k: u64) -> T {
        if k > self.n {
            return T::zero();
        }
        self.ln_pmf(k).exp()
    }

    fn ln_pmf(&self, k: u64) -> T {
        if k > self.n {
            return T::neg_infinity();
        }
        let one = T::one();
        let nf = T::from(self.n).unwrap();
        let kf = T::from(k).unwrap();
        lgamma(nf + one) - lgamma(kf + one) - lgamma(nf - kf + one)
            + kf * self.p.ln()
            + (nf - kf) * (one - self.p).ln()
    }

    fn cdf(&self, k: u64) -> T {
        if k >= self.n {
            return T::one();
        }
        let one = T::one();
        // P(X ≤ k) = I_{1-p}(n-k, k+1)
        let a = T::from(self.n - k).unwrap();
        let b = T::from(k + 1).unwrap();
        betainc(a, b, one - self.p).unwrap_or(T::nan())
    }

    fn quantile(&self, p: T) -> u64 {
        if p <= T::zero() {
            return 0;
        }
        if p >= T::one() {
            return self.n;
        }
        // Normal approximation for initial guess: k0 ≈ mean + std·z_p
        let mean = T::from(self.n).unwrap() * self.p;
        let std = (mean * (T::one() - self.p)).sqrt();
        let z = super::normal_quantile_standard(p);
        let k0 = (mean + std * z)
            .max(T::zero())
            .min(T::from(self.n).unwrap())
            .floor()
            .to_u64()
            .unwrap_or(0);
        super::discrete_quantile_search(|k| self.cdf(k), p, k0)
    }

    fn mean(&self) -> T {
        T::from(self.n).unwrap() * self.p
    }

    fn variance(&self) -> T {
        T::from(self.n).unwrap() * self.p * (T::one() - self.p)
    }
}