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

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

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

/// # The Hypergeometric Distribution
///
/// ## Description
/// Density, distribution function, quantile function and random generation for the hypergeometric
/// distribution.
///
/// ## Arguments
///
/// * Quantiles representing the number of white balls drawn without replacement
/// from an urn which contains both black and white balls.
/// * m: the number of white balls in the urn.
/// * n: the number of black balls in the urn.
/// * k: the number of balls drawn from the urn.
/// * p: probability, it must be between 0 and 1.
///
/// ## Details
///
/// The hypergeometric distribution is used for sampling without replacement. The density of this
/// distribution with parameters m, n and k (named Np, N-Np, and n, respectively in the reference below) is given by
///
/// $ p(x) = {m \choose x} {n \choose k-x} / {m+n \choose k} $
///
/// for x = 0, …, k.
///
/// Note that p(x) is non-zero only for max(0, k-n) <= x <= min(k, m).
///
/// With $p := \frac{m}{m+n}$ (hence $Np = N \times p$ in the reference's notation), the first two moments
/// are mean
///
/// $ E\[X\] = μ = k p $
///
/// and variance
///
/// $ Var(X) = k p (1 - p) * \frac{m+n-k}{m+n-1} $,
///
/// which shows the closeness to the Binomial(k,p) (where the hypergeometric has smaller variance
/// unless k = 1).
///
/// The quantile is defined as the smallest value x such that F(x) ≥ p, where F is the distribution
/// function.
///
/// If one of m, n, k, exceeds .Machine$integer.max, currently the equivalent of
/// qhyper(runif(nn), m,n,k) is used, when a binomial approximation may be considerably
/// more efficient.
///
/// ## Density Plot
///
/// ```rust
/// # use r2rs_base::traits::StatisticalSlice;
/// # use r2rs_nmath::{distribution::HyperGeometricBuilder, traits::Distribution};
/// # use strafe_plot::prelude::{IntoDrawingArea, Line, Plot, PlotOptions, SVGBackend, BLACK};
/// # use strafe_type::FloatConstraint;
/// let hyper_geom = HyperGeometricBuilder::new().build().unwrap();
/// let x = <[f64]>::sequence_by(-0.5, 1.5, 0.001);
/// let y = x
///     .iter()
///     .map(|x| hyper_geom.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/hyper/doctest_out/density.svg"),
/// #     )
/// #     .unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = embed_doc_image::embed_image!("density", "src/distribution/hyper/doctest_out/density.svg")))]
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = "![Density][density]"))]
///
/// ## Source
///
/// dhyper computes via binomial probabilities, using code contributed by Catherine Loader
/// (see dbinom).
///
/// phyper is based on calculating dhyper and phyper(...)/dhyper(...) (as a summation), based on
/// ideas of Ian Smith and Morten Welinder.
///
/// qhyper is based on inversion.
///
/// rhyper is based on a corrected version of
///
/// Kachitvichyanukul, V. and Schmeiser, B. (1985). Computer generation of hypergeometric random
/// variates. Journal of Statistical Computation and Simulation, 22, 127–145.
///
/// ## References
///
/// Johnson, N. L., Kotz, S., and Kemp, A. W. (1992) Univariate Discrete Distributions, Second
/// Edition. New York: Wiley.
///
/// ## See Also
///
/// Distributions for other standard distributions.
///
/// ## Examples
///
/// These are not equal, but the error is very small
/// ```rust
/// # use r2rs_nmath::{distribution::HyperGeometricBuilder, traits::Distribution};
/// # use r2rs_stats::traits::StatArray;
/// # use strafe_type::FloatConstraint;
/// let m = 10;
/// let n = 7;
/// let k = 8;
///
/// let x = (0..=k + 1).collect::<Vec<_>>();
///
/// let hyper = HyperGeometricBuilder::new()
///     .with_group_1(m)
///     .with_group_2(n)
///     .with_number_drawn(k)
///     .build()
///     .unwrap();
/// let p = x
///     .iter()
///     .map(|x| hyper.probability(x, true).unwrap())
///     .collect::<Vec<_>>();
/// let d = x
///     .iter()
///     .map(|x| hyper.density(x).unwrap())
///     .collect::<Vec<_>>()
///     .cumsum();
/// let diff = p
///     .iter()
///     .zip(d.iter())
///     .map(|(p, d)| p - d)
///     .collect::<Vec<_>>();
///
/// println!("{p:?}");
/// println!("{d:?}");
/// println!("{diff:?}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/hyper/doctest_out/difference.md").unwrap();
/// # writeln!(f, "```output").unwrap();
/// # writeln!(f, "{p:?}").unwrap();
/// # writeln!(f, "{d:?}").unwrap();
/// # writeln!(f, "{diff:?}").unwrap();
/// # writeln!(f, "```").unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = include_str!("doctest_out/difference.md")))]
pub struct HyperGeometric {
    group_1: PositiveInteger64,
    group_2: PositiveInteger64,
    number_drawn: PositiveInteger64,
}

impl Distribution for HyperGeometric {
    fn density<R: Into<Real64>>(&self, x: R) -> Real64 {
        dhyper(x, self.group_1, self.group_2, self.number_drawn, false)
    }

    fn log_density<R: Into<Real64>>(&self, x: R) -> Real64 {
        dhyper(x, self.group_1, self.group_2, self.number_drawn, true)
    }

    fn probability<R: Into<Real64>>(&self, q: R, lower_tail: bool) -> Probability64 {
        phyper(q, self.group_1, self.group_2, self.number_drawn, lower_tail)
    }

    fn log_probability<R: Into<Real64>>(&self, q: R, lower_tail: bool) -> LogProbability64 {
        log_phyper(q, self.group_1, self.group_2, self.number_drawn, lower_tail)
    }

    fn quantile<P: Into<Probability64>>(&self, p: P, lower_tail: bool) -> Real64 {
        qhyper(p, self.group_1, self.group_2, self.number_drawn, lower_tail)
    }

    fn log_quantile<LP: Into<LogProbability64>>(&self, p: LP, lower_tail: bool) -> Real64 {
        log_qhyper(p, self.group_1, self.group_2, self.number_drawn, lower_tail)
    }

    fn random_sample<R: RNG>(&self, rng: &mut R) -> Real64 {
        rhyper(self.group_1, self.group_2, self.number_drawn, rng)
    }
}

pub struct HyperGeometricBuilder {
    group_1: Option<PositiveInteger64>,
    group_2: Option<PositiveInteger64>,
    number_drawn: Option<PositiveInteger64>,
}

impl HyperGeometricBuilder {
    pub fn new() -> Self {
        Self {
            group_1: None,
            group_2: None,
            number_drawn: None,
        }
    }

    pub fn with_group_1<P: Into<PositiveInteger64>>(&mut self, group_1: P) -> &mut Self {
        self.group_1 = Some(group_1.into());
        self
    }

    pub fn with_group_2<P: Into<PositiveInteger64>>(&mut self, group_2: P) -> &mut Self {
        self.group_2 = Some(group_2.into());
        self
    }

    pub fn with_number_drawn<P: Into<PositiveInteger64>>(&mut self, number_drawn: P) -> &mut Self {
        self.number_drawn = Some(number_drawn.into());
        self
    }

    pub fn build(&self) -> Result<HyperGeometric, String> {
        let group_1 = self.group_1.unwrap_or(1.0.into());
        let group_2 = self.group_2.unwrap_or(1.0.into());
        let number_drawn = self.number_drawn.unwrap_or(1.0.into());

        if number_drawn.unwrap() > group_1.unwrap() + group_2.unwrap() {
            Err(format!(
                "Number drawn must be less than the two group sizes combined: {} > {}",
                number_drawn.unwrap(),
                group_1.unwrap() + group_2.unwrap()
            ))
        } else {
            Ok(HyperGeometric {
                group_1,
                group_2,
                number_drawn,
            })
        }
    }
}

#[cfg(test)]
mod tests;

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

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