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

/// 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,
    lgamma_n1: T, // lgamma(n + 1), cached at construction
    ln_p: T,      // ln(p), cached at construction
    ln_q: T,      // ln(1 − p), cached at construction
}

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);
        }
        let one = T::one();
        Ok(Self {
            n,
            p,
            lgamma_n1: lgamma(T::from(n).unwrap() + one),
            ln_p: p.ln(),
            ln_q: (one - p).ln(),
        })
    }
}

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)
        }
    }

    impl_sample_array!(u64, 0u64);
}

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();
        self.lgamma_n1 - lgamma(kf + one) - lgamma(nf - kf + one)
            + kf * self.ln_p
            + (nf - kf) * self.ln_q
    }

    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)
    }
}
