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

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

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

/// # The Gamma Distribution
///
/// ## Description:
///
/// Density, distribution function, quantile function and random
/// generation for the Gamma distribution with parameters ‘shape’ and
/// ‘scale’.
///
/// ## Arguments:
///
/// * rate: an alternative way to specify the scale.
/// * shape, scale: shape and scale parameters.  Must be positive, ‘scale’
/// strictly.
///
/// ## Details:
///
/// If ‘scale’ is omitted, it assumes the default value of ‘1’.
///
/// The Gamma distribution with parameters ‘shape’ = a and ‘scale’ = s
/// has density
///
/// $ f(x)= \frac{1}{s^a \Gamma(a)} x^{a-1} e^{-\frac{x}{s}} $
///
/// for $ x >= 0 $, $ a > 0 $ and $ s > 0 $.  (Here Gamma(a) is the function
/// implemented by R's ‘gamma()’ and defined in its help.  Note that a
/// = 0 corresponds to the trivial distribution with all mass at point
/// 0.)
///
/// The mean and variance are $ E(X) = a\*s $ and $ Var(X) = a\*s^2 $.
///
/// The cumulative hazard $H(t) = - log(1 - F(t))$ is
///
/// -pgamma(t, ..., lower = FALSE, log = TRUE)
///
/// Note that for smallish values of ‘shape’ (and moderate ‘scale’) a
/// large parts of the mass of the Gamma distribution is on values of
/// x so near zero that they will be represented as zero in computer
/// arithmetic.  So ‘rgamma’ may well return values which will be
/// represented as zero.  (This will also happen for very large values
/// of ‘scale’ since the actual generation is done for ‘scale = 1’.)
///
/// ## Density Plot
///
/// ```rust
/// # use r2rs_base::traits::StatisticalSlice;
/// # use r2rs_nmath::{distribution::GammaBuilder, traits::Distribution};
/// # use strafe_plot::prelude::{IntoDrawingArea, Line, Plot, PlotOptions, SVGBackend, BLACK};
/// # use strafe_type::FloatConstraint;
/// let gamma = GammaBuilder::new().build();
/// let x = <[f64]>::sequence(-0.5, 4.0, 1000);
/// let y = x
///     .iter()
///     .map(|x| gamma.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/gamma/doctest_out/density.svg"),
/// #     )
/// #     .unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = embed_doc_image::embed_image!("density", "src/distribution/gamma/doctest_out/density.svg")))]
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = "![Density][density]"))]
///
/// ## Note:
///
/// The S (Becker _et al_, 1988) parametrization was via ‘shape’ and
/// ‘rate’: S had no ‘scale’ parameter. It is an error to supply and
/// ‘scale’ and ‘rate’.
///
/// ‘pgamma’ is closely related to the incomplete gamma function.  As
/// defined by Abramowitz and Stegun 6.5.1 (and by ‘Numerical
/// Recipes’) this is
///
/// $ P(a,x) = 1/\Gamma(a) \int_0^x t^{a-1} exp(-t) dt $
///
/// $P(a, x)$ is ‘pgamma(x, a)’.  Other authors (for example Karl
/// Pearson in his 1922 tables) omit the normalizing factor, defining
/// the incomplete gamma function gamma(a,x) as
/// $ gamma(a,x) = \int_0^x t^{a-1} exp(-t) dt $, i.e.,
/// ‘$ pgamma(x, a) * gamma(a) $’.
/// Yet other use the ‘upper’ incomplete gamma function,
///
/// $ Gamma(a,x) = \int_x^\infty t^{a-1} exp(-t) dt $,
///
/// which can be computed by ‘pgamma(x, a, lower = FALSE) * gamma(a)’.
///
/// Note however that ‘pgamma(x, a, ..)’ currently requires a > 0,
/// whereas the incomplete gamma function is also defined for negative
/// a.  In that case, you can use ‘gamma_inc(a,x)’ (for Gamma(a,x))
/// from package ‘gsl’.
///
/// See also <URL:
/// https:///en.wikipedia.org/wiki/Incomplete_gamma_function>, or <URL:
/// http:///dlmf.nist.gov/8.2#i>.
///
/// ## Source:
///
/// ‘dgamma’ is computed via the Poisson density, using code
/// contributed by Catherine Loader (see ‘dbinom’).
///
/// ‘pgamma’ uses an unpublished (and not otherwise documented)
/// algorithm ‘mainly by Morten Welinder’.
///
/// ‘qgamma’ is based on a C translation of
///
/// Best, D. J. and D. E. Roberts (1975).  Algorithm AS91. Percentage
/// points of the chi-squared distribution.  _Applied Statistics_,
/// *24*, 385-388.
///
/// plus a final Newton step to improve the approximation.
///
/// ‘rgamma’ for ‘shape >= 1’ uses
///
/// Ahrens, J. H. and Dieter, U. (1982).  Generating gamma variates by
/// a modified rejection technique.  _Communications of the ACM_,
/// *25*, 47-54,
///
/// and for ‘0 < shape < 1’ uses
///
/// Ahrens, J. H. and Dieter, U. (1974).  Computer methods for
/// sampling from gamma, beta, Poisson and binomial distributions.
/// _Computing_, *12*, 223-246.
///
/// ## References:
///
/// Becker, R. A., Chambers, J. M. and Wilks, A. R. (1988).  _The New
/// S Language_.  Wadsworth & Brooks/Cole.
///
/// Shea, B. L. (1988).  Algorithm AS 239: Chi-squared and incomplete
/// Gamma integral, _Applied Statistics (JRSS C)_, *37*, 466-473.
/// doi: 10.2307/2347328 (URL: https:///doi.org/10.2307/2347328).
///
/// Abramowitz, M. and Stegun, I. A. (1972) _Handbook of Mathematical
/// Functions._ New York: Dover.  Chapter 6: Gamma and Related
/// Functions.
///
/// NIST Digital Library of Mathematical Functions.  <URL:
/// http:///dlmf.nist.gov/>, section 8.2.
///
/// ## See Also:
///
/// ‘gamma’ for the gamma function.
///
/// Distributions for other standard distributions, including ‘dbeta’
/// for the Beta distribution and ‘dchisq’ for the chi-squared
/// distribution which is a special case of the Gamma distribution.
///
/// ## Examples:
///
/// ```rust
/// # use r2rs_nmath::{distribution::GammaBuilder, traits::Distribution};
/// # use strafe_type::FloatConstraint;
/// let gamma = GammaBuilder::new().with_shape(1).build();
/// let r = (1..=4)
///     .map(|x| -gamma.density(x).unwrap().ln())
///     .collect::<Vec<_>>();
/// println!("{r:?}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/gamma/doctest_out/dens.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/dens.md")))]
///
/// ```rust
/// # use r2rs_nmath::{distribution::GammaBuilder, traits::Distribution};
/// # use strafe_type::FloatConstraint;
/// let p = (1..=9).map(|i| i as f64 / 10.0).collect::<Vec<_>>();
/// let gamma = GammaBuilder::new().with_shape(2).build();
/// let r = p
///     .iter()
///     .map(|p| gamma.probability(gamma.quantile(p, true), true).unwrap())
///     .collect::<Vec<_>>();
/// println!("{r:?}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/gamma/doctest_out/prob1.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/prob1.md")))]
///
/// ```rust
/// # use r2rs_nmath::{distribution::GammaBuilder, traits::Distribution};
/// # use strafe_type::FloatConstraint;
/// let p = (1..=9).map(|i| i as f64 / 10.0).collect::<Vec<_>>();
/// let gamma = GammaBuilder::new().with_shape(1).build();
/// let r = p
///     .iter()
///     .map(|p| 1.0 - 1.0 / gamma.quantile(p, true).unwrap().exp())
///     .collect::<Vec<_>>();
/// println!("{r:?}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/gamma/doctest_out/prob2.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/prob2.md")))]
///
/// Even for shape = 0.001 about half the mass is on numbers
/// that cannot be represented accurately (and most of those as zero)
/// ```rust
/// # use num_traits::Float;
/// # use r2rs_nmath::{
/// #     distribution::GammaBuilder,
/// #     rng::MersenneTwister,
/// #     traits::{Distribution, RNG},
/// # };
/// # use strafe_type::FloatConstraint;
/// let gamma = GammaBuilder::new().with_shape(0.001).build();
/// let r1 = gamma.probability(f64::min_positive_value(), true);
/// println!("{r1}");
/// let r2 = gamma.probability(5e-324, true);
/// println!("{r2}");
/// let mut rng = MersenneTwister::new();
/// rng.set_seed(1);
/// let r3 = (0..10_000)
///     .map(|_| {
///         if gamma.random_sample(&mut rng).unwrap() == 0.0 {
///             1.0
///         } else {
///             0.0
///         }
///     })
///     .sum::<f64>();
/// println!(
///     "{}% of the samples are 0",
///     ((r3 / 10_000.0) * 100.0) as usize
/// );
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/gamma/doctest_out/limits.md").unwrap();
/// # writeln!(f, "```output").unwrap();
/// # writeln!(f, "{r1}").unwrap();
/// # writeln!(f, "{r2}").unwrap();
/// # writeln!(f, "{}% of the samples are 0", ((r3 / 10_000.0) * 100.0) as usize).unwrap();
/// # writeln!(f, "```").unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = include_str!("doctest_out/limits.md")))]
pub struct Gamma {
    shape: Positive64,
    scale: Rational64,
}

impl Distribution for Gamma {
    fn density<R: Into<Real64>>(&self, x: R) -> Real64 {
        dgamma(x, self.shape, self.scale, false)
    }

    fn log_density<R: Into<Real64>>(&self, x: R) -> Real64 {
        dgamma(x, self.shape, self.scale, true)
    }

    fn probability<R: Into<Real64>>(&self, q: R, lower_tail: bool) -> Probability64 {
        pgamma(q, self.shape, self.scale, lower_tail)
    }

    fn log_probability<R: Into<Real64>>(&self, q: R, lower_tail: bool) -> LogProbability64 {
        log_pgamma(q, self.shape, self.scale, lower_tail)
    }

    fn quantile<P: Into<Probability64>>(&self, p: P, lower_tail: bool) -> Real64 {
        qgamma(p, self.shape, self.scale, lower_tail)
    }

    fn log_quantile<LP: Into<LogProbability64>>(&self, p: LP, lower_tail: bool) -> Real64 {
        log_qgamma(p, self.shape, self.scale, lower_tail)
    }

    fn random_sample<R: RNG>(&self, rng: &mut R) -> Real64 {
        rgamma(self.shape, self.scale, rng)
    }
}

pub struct GammaBuilder {
    shape: Option<Positive64>,
    scale: Option<Rational64>,
}

impl GammaBuilder {
    pub fn new() -> Self {
        Self {
            shape: None,
            scale: None,
        }
    }

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

    pub fn with_scale<R: Into<Rational64>>(&mut self, scale: R) -> &mut Self {
        self.scale = Some(scale.into());
        self
    }

    pub fn build(&self) -> Gamma {
        let shape = self.shape.unwrap_or(1.0.into());
        let scale = self.scale.unwrap_or(1.0.into());

        Gamma { shape, scale }
    }
}

#[cfg(test)]
mod tests;

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

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