mod d;
mod p;
mod q;
mod q_mu;
mod r;

use strafe_type::{LogProbability64, Positive64, Probability64, Real64};

pub use self::{d::*, p::*, q::*, q_mu::*, r::*};
use crate::traits::{Distribution, RNG};

/// # The Negative Binomial Distribution
///
/// ## Description
/// Density, distribution function, quantile function and random generation for the negative
/// binomial distribution with parameters size and prob.
///
/// ## Arguments
/// * size: target for number of successful trials, or dispersion parameter (the shape parameter of
/// the gamma mixing distribution). Must be strictly positive, need not be integer.
/// * prob: probability of success in each trial. 0 < prob <= 1.
/// * mu: alternative parametrization via mean: see ‘Details’.
///
/// ## Details
///
/// The negative binomial distribution with size = n and prob = p has density
///
/// $ \frac{\Gamma(x+n)}{(\Gamma(n) x!)} p^n (1-p)^x $
///
/// for $ x = 0, 1, 2, …, n > 0 $ and $ 0 < p ≤ 1 $.
///
/// This represents the number of failures which occur in a sequence of Bernoulli trials before
/// a target number of successes is reached. The mean is $ \mu = n\frac{1-p}{p} $ and variance
/// $ n\frac{1-p}{p^2} $.
///
/// A negative binomial distribution can also arise as a mixture of Poisson distributions with
/// mean distributed as a gamma distribution (see pgamma) with scale parameter
/// $ \frac{1 - prob}{prob} $ and shape parameter size. (This definition allows non-integer values
/// of size.)
///
/// An alternative parametrization (often used in ecology) is by the mean mu (see above), and
/// size, the dispersion parameter, where $ prob = \frac{\text{size}}{\text{size}+\mu} $. The
/// variance is $ \mu + \frac{\mu^2}{\text{size}} $ in this parametrization.
///
/// If an element of x is not integer, the result of dnbinom is zero, with a warning.
///
/// The case size == 0 is the distribution concentrated at zero. This is the limiting distribution
/// for size approaching zero, even if mu rather than prob is held constant. Notice though, that
/// the mean of the limit distribution is 0, whatever the value of $ \mu $.
///
/// The quantile is defined as the smallest value x such that F(x) ≥ p, where F is the distribution
/// function.
///
/// ## Density Plot
///
/// ```rust
/// # use r2rs_base::traits::StatisticalSlice;
/// # use r2rs_nmath::{distribution::NegativeBinomialBuilder, traits::Distribution};
/// # use strafe_plot::prelude::{IntoDrawingArea, Line, Plot, PlotOptions, SVGBackend, BLACK};
/// # use strafe_type::FloatConstraint;
/// let nbeta = NegativeBinomialBuilder::new().build();
/// let x = <[f64]>::sequence_by(-1.0, 1.0, 0.001);
/// let y = x
///     .iter()
///     .map(|x| nbeta.density(x).unwrap())
///     .collect::<Vec<_>>();
///
/// let root = SVGBackend::new("density.svg", (1024, 768)).into_drawing_area();
/// Plot::new()
///     .with_options(PlotOptions {
///         x_axis_label: "x".to_string(),
///         y_axis_label: "density".to_string(),
///         ..Default::default()
///     })
///     .with_plottable(Line {
///         x,
///         y,
///         color: BLACK,
///         ..Default::default()
///     })
///     .plot(&root)
///     .unwrap();
/// # use std::fs::rename;
/// #     drop(root);
/// #     rename(
/// #             format!("density.svg"),
/// #             format!("src/distribution/nbinom/doctest_out/density.svg"),
/// #     )
/// #     .unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = embed_doc_image::embed_image!("density", "src/distribution/nbinom/doctest_out/density.svg")))]
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = "![Density][density]"))]
///
/// ## Source
///
/// dnbinom computes via binomial probabilities, using code contributed by Catherine Loader
/// (see dbinom).
///
/// pnbinom uses pbeta.
///
/// qnbinom uses the Cornish–Fisher Expansion to include a skewness correction to a normal
/// approximation, followed by a search.
///
/// rnbinom uses the derivation as a gamma mixture of Poissons, see
///
/// Devroye, L. (1986) Non-Uniform Random Variate Generation. Springer-Verlag, New York. Page 480.
///
/// ## See Also
///
/// Distributions for standard distributions, including dbinom for the binomial, dpois for the
/// Poisson and dgeom for the geometric distribution, which is a special case of the negative
/// binomial.
///
/// ## Examples
///
/// Equals one
/// ```rust
/// # use r2rs_nmath::{distribution::NegativeBinomialBuilder, traits::Distribution};
/// # use strafe_type::FloatConstraint;
/// let x = (0..=11).collect::<Vec<_>>();
/// let nbinom = NegativeBinomialBuilder::new()
///     .with_size(1)
///     .with_success_probability(0.5)
///     .build();
/// let r = x
///     .iter()
///     .map(|x| nbinom.density(x).unwrap() * 2.0_f64.powi(1 + x))
///     .collect::<Vec<_>>();
/// println!("{r:?}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/nbinom/doctest_out/equals_one.md").unwrap();
/// # writeln!(f, "```output").unwrap();
/// # writeln!(f, "{r:?}").unwrap();
/// # writeln!(f, "```").unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = include_str!("doctest_out/equals_one.md")))]
///
/// Theoretically integer
/// ```rust
/// # use r2rs_nmath::{distribution::NegativeBinomialBuilder, traits::Distribution};
/// # use strafe_type::FloatConstraint;
/// let x = (0..=8).collect::<Vec<_>>();
/// let nbinom = NegativeBinomialBuilder::new()
///     .with_size(2)
///     .with_success_probability(0.5)
///     .build();
/// let r = x
///     .iter()
///     .map(|x| 126.0 / nbinom.density(x).unwrap())
///     .collect::<Vec<_>>();
/// println!("{r:?}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/nbinom/doctest_out/integer.md").unwrap();
/// # writeln!(f, "```output").unwrap();
/// # writeln!(f, "{r:?}").unwrap();
/// # writeln!(f, "```").unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = include_str!("doctest_out/integer.md")))]
///
/// Cumulative ('p') = Sum of discrete prob.s ('d')
/// ```rust
/// # use r2rs_nmath::{distribution::NegativeBinomialBuilder, traits::Distribution};
/// # use r2rs_stats::traits::StatArray;
/// # use strafe_type::FloatConstraint;
/// let x = (0..=11).collect::<Vec<_>>();
/// let nbinom = NegativeBinomialBuilder::new()
///     .with_size(2)
///     .with_success_probability(0.5)
///     .build();
/// let d_sum = x
///     .iter()
///     .map(|x| nbinom.density(x).unwrap())
///     .collect::<Vec<_>>()
///     .cumsum();
/// let p = x
///     .iter()
///     .map(|x| nbinom.probability(x, true).unwrap())
///     .collect::<Vec<_>>();
/// let r = d_sum
///     .iter()
///     .zip(p.iter())
///     .map(|(d, p)| 1.0 - d / p)
///     .collect::<Vec<_>>();
/// println!("{r:?}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/nbinom/doctest_out/sum_dens_prob.md").unwrap();
/// # writeln!(f, "```output").unwrap();
/// # writeln!(f, "{r:?}").unwrap();
/// # writeln!(f, "```").unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = include_str!("doctest_out/sum_dens_prob.md")))]
///
/// ```r
/// x <- 0:15
/// size <- (1:20)/4
/// persp(x, size, dnb <- outer(x, size, function(x,s) dnbinom(x, s, prob = 0.4)),
///       xlab = "x", ylab = "s", zlab = "density", theta = 150)
/// title(tit <- "negative binomial density(x,s, pr = 0.4)  vs.  x & s")
///
/// image  (x, size, log10(dnb), main = paste("log [", tit, "]"))
/// contour(x, size, log10(dnb), add = TRUE)
///
/// ## Alternative parametrization
/// x1 <- rnbinom(500, mu = 4, size = 1)
/// x2 <- rnbinom(500, mu = 4, size = 10)
/// x3 <- rnbinom(500, mu = 4, size = 100)
/// h1 <- hist(x1, breaks = 20, plot = FALSE)
/// h2 <- hist(x2, breaks = h1$breaks, plot = FALSE)
/// h3 <- hist(x3, breaks = h1$breaks, plot = FALSE)
/// barplot(rbind(h1$counts, h2$counts, h3$counts),
///         beside = TRUE, col = c("red","blue","cyan"),
///         names.arg = round(h1$breaks[-length(h1$breaks)]))
/// ```
pub struct NegativeBinomial {
    success_probability: Probability64,
    size: Positive64,
    mu: Option<Positive64>,
}

impl Distribution for NegativeBinomial {
    fn density<R: Into<Real64>>(&self, x: R) -> Real64 {
        if let Some(mu) = self.mu {
            dnbinom_mu(x, self.size, mu, false)
        } else {
            dnbinom(x, self.size, self.success_probability, false)
        }
    }

    fn log_density<R: Into<Real64>>(&self, x: R) -> Real64 {
        if let Some(mu) = self.mu {
            dnbinom_mu(x, self.size, mu, true)
        } else {
            dnbinom(x, self.size, self.success_probability, true)
        }
    }

    fn probability<R: Into<Real64>>(&self, q: R, lower_tail: bool) -> Probability64 {
        if let Some(mu) = self.mu {
            pnbinom_mu(q, self.size, mu, lower_tail)
        } else {
            pnbinom(q, self.size, self.success_probability, lower_tail)
        }
    }

    fn log_probability<R: Into<Real64>>(&self, q: R, lower_tail: bool) -> LogProbability64 {
        if let Some(mu) = self.mu {
            log_pnbinom_mu(q, self.size, mu, lower_tail)
        } else {
            log_pnbinom(q, self.size, self.success_probability, lower_tail)
        }
    }

    fn quantile<P: Into<Probability64>>(&self, p: P, lower_tail: bool) -> Real64 {
        if let Some(mu) = self.mu {
            qnbinom_mu(p, self.size, mu, lower_tail)
        } else {
            qnbinom(p, self.size, self.success_probability, lower_tail)
        }
    }

    fn log_quantile<LP: Into<LogProbability64>>(&self, p: LP, lower_tail: bool) -> Real64 {
        if let Some(mu) = self.mu {
            log_qnbinom_mu(p, self.size, mu, lower_tail)
        } else {
            log_qnbinom(p, self.size, self.success_probability, lower_tail)
        }
    }

    fn random_sample<R: RNG>(&self, rng: &mut R) -> Real64 {
        if let Some(mu) = self.mu {
            rnbinom_mu(self.size, mu, rng)
        } else {
            rnbinom(self.size, self.success_probability, rng)
        }
    }
}

pub struct NegativeBinomialBuilder {
    success_probability: Option<Probability64>,
    size: Option<Positive64>,
    mu: Option<Positive64>,
}

impl NegativeBinomialBuilder {
    pub fn new() -> Self {
        Self {
            size: None,
            success_probability: None,
            mu: None,
        }
    }

    pub fn with_size<P: Into<Positive64>>(&mut self, size: P) -> &mut Self {
        self.size = Some(size.into());
        self
    }

    pub fn with_success_probability<P: Into<Probability64>>(
        &mut self,
        success_probability: P,
    ) -> &mut Self {
        self.success_probability = Some(success_probability.into());
        self
    }

    pub fn with_mu<P: Into<Positive64>>(&mut self, mu: P) -> &mut Self {
        self.mu = Some(mu.into());
        self
    }

    pub fn build(&self) -> NegativeBinomial {
        let success_probability = self.success_probability.unwrap_or(1.0.into());
        let size = self.size.unwrap_or(1.0.into());

        NegativeBinomial {
            size,
            success_probability,
            mu: self.mu,
        }
    }
}

#[cfg(test)]
mod tests;

#[cfg(all(test, feature = "enable_proptest"))]
mod proptests;

#[cfg(all(test, feature = "enable_covtest"))]
mod covtests;