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

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

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

/// # The Binomial Distribution
///
/// ## Description:
///
/// Density, distribution function, quantile function and random
/// generation for the binomial distribution with parameters ‘size’
/// and ‘prob’.
///
/// This is conventionally interpreted as the number of ‘successes’ in
/// ‘size’ trials.
///
/// ## Arguments:
///
/// * size: number of trials (zero or more).
/// * prob: probability of success on each trial.
///
/// ## Details:
///
/// The binomial distribution with ‘size’ = n and ‘prob’ = p has
/// density
///
/// $p(x) = choose(n, x) p^x (1-p)^(n-x)$
///
/// for x = 0, ..., n. Note that binomial _coefficients_ can be
/// computed by ‘choose’ in R.
///
/// If an element of ‘x’ is not integer, the result of ‘dbinom’ is
/// zero, with a warning.
///
/// p(x) is computed using Loader's algorithm, see the reference
/// below.
///
/// 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::BinomialBuilder, traits::Distribution};
/// # use strafe_plot::prelude::{IntoDrawingArea, Line, Plot, PlotOptions, SVGBackend, BLACK};
/// # use strafe_type::FloatConstraint;
/// let binom = BinomialBuilder::new().build();
/// let x = <[f64]>::sequence_by(-0.5, 1.5, 0.001);
/// let y = x
///     .iter()
///     .map(|x| binom.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/binom/doctest_out/density.svg"),
/// #     )
/// #     .unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = embed_doc_image::embed_image!("density", "src/distribution/binom/doctest_out/density.svg")))]
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = "![Density][density]"))]
///
/// ## Source:
///
/// For ‘dbinom’ a saddle-point expansion is used: see
///
/// Catherine Loader (2000). _Fast and Accurate Computation of
/// Binomial Probabilities_; available from \<URL:
/// <http://www.herine.net/stat/software/dbinom.html>\>.
///
/// ‘pbinom’ uses ‘pbeta’.
///
/// ‘qbinom’ uses the Cornish-Fisher Expansion to include a skewness
/// correction to a normal approximation, followed by a search.
///
/// ‘rbinom’ (for ‘size < .Machine$integer.max’) is based on
///
/// Kachitvichyanukul, V. and Schmeiser, B. W. (1988) Binomial random
/// variate generation. _Communications of the ACM_, *31*, 216-222.
///
/// For larger values it uses inversion.
///
/// ## See Also:
///
/// Distributions for other standard distributions, including
/// ‘dnbinom’ for the negative binomial, and ‘dpois’ for the Poisson
/// distribution.
///
/// ## Examples:
///
/// Compute $P(45 < X < 55)$ for $Binomial(100,0.5)$
/// ```rust
/// # use r2rs_nmath::{distribution::BinomialBuilder, traits::Distribution};
/// # use strafe_type::FloatConstraint;
/// let binom = BinomialBuilder::new()
///     .with_size(100)
///     .with_success_probability(0.5)
///     .build();
/// let p = (46..=54).map(|x| binom.density(x).unwrap()).sum::<f64>();
/// println!("{p:?}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/binom/doctest_out/binomial_sum.md").unwrap();
/// # writeln!(f, "```output").unwrap();
/// # writeln!(f, "{p:?}").unwrap();
/// # writeln!(f, "```").unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = include_str!("doctest_out/binomial_sum.md")))]
/// ```rust
/// # use std::fs::rename;
/// # use num_traits::FloatConst;
/// # use r2rs_base::traits::StatisticalSlice;
/// # use r2rs_nmath::{distribution::BinomialBuilder, traits::Distribution};
/// # use strafe_plot::prelude::*;
/// # use strafe_type::FloatConstraint;
/// let n = 2000.0;
/// let k = <[f64]>::sequence_by(0.0, n, 20.0);
/// let binom = BinomialBuilder::new()
///     .with_size(n)
///     .with_success_probability(f64::PI() / 10.0)
///     .build();
/// let y1 = k
///     .iter()
///     .map(|k| binom.log_density(k).unwrap())
///     .collect::<Vec<_>>();
/// let y2 = k
///     .iter()
///     .map(|k| binom.density(k).unwrap().ln())
///     .collect::<Vec<_>>();
///
/// Plot::new()
///     .with_options(PlotOptions {
///         title: "log_density() is better than density().ln()".to_string(),
///         x_axis_label: "k".to_string(),
///         y_axis_label: "log density".to_string(),
///         legend_x: 0.9,
///         legend_y: 0.1,
///         ..Default::default()
///     })
///     .with_plottable(Line {
///         x: k.clone(),
///         y: y1,
///         color: BLACK,
///         legend: true,
///         label: "log_dbinom()".to_string(),
///         ..Default::default()
///     })
///     .with_plottable(Line {
///         x: k.clone(),
///         y: y2,
///         color: RED,
///         legend: true,
///         label: "dbinom().ln()".to_string(),
///         ..Default::default()
///     })
///     .plot(&SVGBackend::new("log_density.svg", (1024, 768)).into_drawing_area())
///     .unwrap();
/// # rename(
/// #     format!("log_density.svg"),
/// #     format!("src/distribution/binom/doctest_out/log_density.svg"),
/// # )
/// # .unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = embed_doc_image::embed_image!("log_density", "src/distribution/binom/doctest_out/log_density.svg")))]
#[cfg_attr(
    feature = "doc_outputs",
    cfg_attr(all(), doc = "![Log Density Plot][log_density]")
)]
pub struct Binomial {
    size: PositiveInteger64,
    success_probability: Probability64,
}

impl Distribution for Binomial {
    fn density<R: Into<Real64>>(&self, x: R) -> Real64 {
        dbinom(x, self.size, self.success_probability, false)
    }

    fn log_density<R: Into<Real64>>(&self, x: R) -> Real64 {
        dbinom(x, self.size, self.success_probability, true)
    }

    fn probability<R: Into<Real64>>(&self, q: R, lower_tail: bool) -> Probability64 {
        pbinom(q, self.size, self.success_probability, lower_tail)
    }

    fn log_probability<R: Into<Real64>>(&self, q: R, lower_tail: bool) -> LogProbability64 {
        log_pbinom(q, self.size, self.success_probability, lower_tail)
    }

    fn quantile<P: Into<Probability64>>(&self, p: P, lower_tail: bool) -> Real64 {
        qbinom(p, self.size, self.success_probability, lower_tail)
    }

    fn log_quantile<LP: Into<LogProbability64>>(&self, p: LP, lower_tail: bool) -> Real64 {
        log_qbinom(p, self.size, self.success_probability, lower_tail)
    }

    fn random_sample<R: RNG>(&self, rng: &mut R) -> Real64 {
        rbinom(self.size, self.success_probability, rng)
    }
}

pub struct BinomialBuilder {
    size: Option<PositiveInteger64>,
    success_probability: Option<Probability64>,
}

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

    pub fn with_size<PI: Into<PositiveInteger64>>(&mut self, size: PI) -> &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 build(&self) -> Binomial {
        let size = self.size.unwrap_or(1.into());
        let success_probability = self.success_probability.unwrap_or(0.5.into());

        Binomial {
            size,
            success_probability,
        }
    }
}

#[cfg(test)]
mod tests;

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

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