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

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

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

/// # The Chi-Squared Distribution
///
/// ## Description:
///
/// Density, distribution function, quantile function and random
/// generation for the chi-squared ($chi^2$) distribution with
/// degrees of freedom and non-centrality parameter.
///
/// ## Arguments:
/// * df: degrees of freedom (non-negative, but can be non-integer).
/// * ncp: non-centrality parameter (non-negative).
///
/// ## Details:
///
/// The chi-squared distribution with ‘df’= n >= 0 degrees of freedom
/// has density
///
/// $f_n(x) = \frac{1}{(2^{\frac{n}{2}} \Gamma(\frac{n}{2}))}  x^{\frac{n}{2}-1} e^{-\frac{x}{2}}$
///
/// for $x > 0$.  The mean and variance are n and 2n.
///
/// The non-central chi-squared distribution with ‘df’= n degrees of
/// freedom and non-centrality parameter ‘ncp’ = $\lambda$ has density
///
/// $ f(x) = \text{exp}(-\frac{\lambda}{2}) \sum_{r=0}^{\infty} (\frac{(\lambda / 2)^r}{r!}) \text{dchisq}(x, df + 2r) $
///
/// for $x >= 0$.  For integer n, this is the distribution of the sum of
/// squares of n normals each with variance one, lambda being the sum
/// of squares of the normal means; further,
///
/// $E(X) = n + \lambda$, $Var(X) = 2(n + 2*\lambda)$, and $E((X - E(X))^3) = 8(n + 3*\lambda)$.
///
/// Note that the degrees of freedom ‘df’= n, can be non-integer, and
/// also n = 0 which is relevant for non-centrality lambda > 0, see
/// Johnson _et al_ (1995, chapter 29).  In that (noncentral, zero df)
/// case, the distribution is a mixture of a point mass at x = 0 (of
/// size ‘pchisq(0, df=0, ncp=ncp)’) and a continuous part, and
/// ‘dchisq()’ is _not_ a density with respect to that mixture measure
/// but rather the limit of the density for df -> 0.
///
/// Note that ‘ncp’ values larger than about 1e5 may give inaccurate
/// results with many warnings for ‘pchisq’ and ‘qchisq’.
///
/// ## Density Plot
///
/// ```rust
/// # use r2rs_base::traits::StatisticalSlice;
/// # use r2rs_nmath::{distribution::ChiSquaredBuilder, traits::Distribution};
/// # use strafe_plot::prelude::{IntoDrawingArea, Line, Plot, PlotOptions, SVGBackend, BLACK};
/// # use strafe_type::FloatConstraint;
/// let chisq = ChiSquaredBuilder::new().build();
/// let x = <[f64]>::sequence(-0.5, 4.0, 1000);
/// let y = x
///     .iter()
///     .map(|x| chisq.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/chisq/doctest_out/density.svg"),
/// #     )
/// #     .unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = embed_doc_image::embed_image!("density", "src/distribution/chisq/doctest_out/density.svg")))]
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = "![Density][density]"))]
///
/// ## Note:
///
/// Supplying ‘ncp = 0’ uses the algorithm for the non-central
/// distribution, which is not the same algorithm used if ‘ncp’ is
/// omitted.  This is to give consistent behaviour in extreme cases
/// with values of ‘ncp’ very near zero.
///
/// The code for non-zero ‘ncp’ is principally intended to be used for
/// moderate values of ‘ncp’: it will not be highly accurate,
/// especially in the tails, for large values.
///
/// ## Source:
///
/// The central cases are computed via the gamma distribution.
///
/// The non-central ‘dchisq’ and ‘rchisq’ are computed as a Poisson
/// mixture central of chi-squares (Johnson _et al_, 1995, p.436).
///
/// The non-central ‘pchisq’ is for ‘ncp < 80’ computed from the
/// Poisson mixture of central chi-squares and for larger ‘ncp’ _via_
/// a C translation of
///
/// Ding, C. G. (1992) Algorithm AS275: Computing the non-central
/// chi-squared distribution function. _Appl.Statist._, *41* 478-482.
///
/// which computes the lower tail only (so the upper tail suffers from
/// cancellation and a warning will be given when this is likely to be
/// significant).
///
/// The non-central ‘qchisq’ is based on inversion of ‘pchisq’.
///
/// ## References:
///
/// Becker, R. A., Chambers, J. M. and Wilks, A. R. (1988) _The New S
/// Language_.  Wadsworth & Brooks/Cole.
///
/// Johnson, N. L., Kotz, S. and Balakrishnan, N. (1995) _Continuous
/// Univariate Distributions_, chapters 18 (volume 1) and 29 (volume
/// 2). Wiley, New York.
///
/// ## See Also:
///
/// A central chi-squared distribution with n degrees of freedom is
/// the same as a Gamma distribution with ‘shape’ a = n/2 and ‘scale’
/// s = 2.
///
/// ## Examples:
///
/// ```rust
/// # use r2rs_nmath::{distribution::ChiSquaredBuilder, traits::Distribution};
/// let d1 = ChiSquaredBuilder::new().with_df(1).build().density(1);
/// let d2 = ChiSquaredBuilder::new().with_df(2).build().density(1);
/// let d3 = ChiSquaredBuilder::new().with_df(3).build().density(1);
///
/// println!("{d1}");
/// println!("{d2}");
/// println!("{d3}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/chisq/doctest_out/dens_df1_3.md").unwrap();
/// # writeln!(f, "```output").unwrap();
/// # writeln!(f, "{d1}").unwrap();
/// # writeln!(f, "{d2}").unwrap();
/// # writeln!(f, "{d3}").unwrap();
/// # writeln!(f, "```").unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = include_str!("doctest_out/dens_df1_3.md")))]
///
/// ```rust
/// # use r2rs_nmath::{distribution::ChiSquaredBuilder, traits::Distribution};
/// let p1 = ChiSquaredBuilder::new()
///     .with_df(3)
///     .build()
///     .probability(1, true);
///
/// println!("{p1}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/chisq/doctest_out/prob_df3.md").unwrap();
/// # writeln!(f, "```output").unwrap();
/// # writeln!(f, "{p1}").unwrap();
/// # writeln!(f, "```").unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = include_str!("doctest_out/prob_df3.md")))]
///
/// Includes the above
/// ```rust
/// # use r2rs_nmath::{distribution::ChiSquaredBuilder, traits::Distribution};
/// let np0 = ChiSquaredBuilder::new()
///     .with_df(3)
///     .with_ncp(0)
///     .build()
///     .probability(1, true);
/// let np1 = ChiSquaredBuilder::new()
///     .with_df(3)
///     .with_ncp(1)
///     .build()
///     .probability(1, true);
/// let np2 = ChiSquaredBuilder::new()
///     .with_df(3)
///     .with_ncp(2)
///     .build()
///     .probability(1, true);
/// let np3 = ChiSquaredBuilder::new()
///     .with_df(3)
///     .with_ncp(3)
///     .build()
///     .probability(1, true);
/// let np4 = ChiSquaredBuilder::new()
///     .with_df(3)
///     .with_ncp(4)
///     .build()
///     .probability(1, true);
///
/// println!("{np0}");
/// println!("{np1}");
/// println!("{np2}");
/// println!("{np3}");
/// println!("{np4}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/chisq/doctest_out/prob_np0_4.md").unwrap();
/// # writeln!(f, "```output").unwrap();
/// # writeln!(f, "{np0}").unwrap();
/// # writeln!(f, "{np1}").unwrap();
/// # writeln!(f, "{np2}").unwrap();
/// # writeln!(f, "{np3}").unwrap();
/// # writeln!(f, "{np4}").unwrap();
/// # writeln!(f, "```").unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = include_str!("doctest_out/prob_np0_4.md")))]
///
/// Non-central RNG -- df = 0 with ncp > 0:  Z0 has point mass at 0
/// ```rust
/// # use r2rs_nmath::{
/// #     distribution::ChiSquaredBuilder,
/// #     rng::MersenneTwister,
/// #     traits::{Distribution, RNG},
/// # };
/// # use strafe_plot::prelude::stem_leaf_plot;
/// # use strafe_type::FloatConstraint;
/// let mut rng = MersenneTwister::new();
/// rng.set_seed(1);
/// let cauchy = ChiSquaredBuilder::new().with_df(1e-100).with_ncp(2).build();
/// let z0 = (0..100)
///     .map(|_| cauchy.random_sample(&mut rng).unwrap())
///     .collect::<Vec<_>>();
/// let slp = stem_leaf_plot(&z0);
/// println!("{slp}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/chisq/doctest_out/stem_leaf.md").unwrap();
/// # writeln!(f, "```output").unwrap();
/// # writeln!(f, "{slp}").unwrap();
/// # writeln!(f, "```").unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = include_str!("doctest_out/stem_leaf.md")))]
///
/// Chi-squared(df = 2) is a special exponential distribution
/// ```rust
/// # use r2rs_nmath::{
/// #     distribution::{ChiSquaredBuilder, ExponentialBuilder},
/// #     traits::Distribution,
/// # };
/// # use strafe_type::FloatConstraint;
/// let chisq = ChiSquaredBuilder::new().with_df(2).build();
/// let exp = ExponentialBuilder::new().with_rate(0.5).build();
/// let chisq_dens = (1..=10)
///     .map(|x| chisq.density(x).unwrap())
///     .collect::<Vec<_>>();
/// println!("{chisq_dens:?}");
/// let exp_dens = (1..=10)
///     .map(|x| exp.density(x).unwrap())
///     .collect::<Vec<_>>();
/// println!("{exp_dens:?}");
/// let chisq_prob = (1..=10)
///     .map(|x| chisq.probability(x, true).unwrap())
///     .collect::<Vec<_>>();
/// println!("{chisq_prob:?}");
/// let exp_prob = (1..=10)
///     .map(|x| exp.probability(x, true).unwrap())
///     .collect::<Vec<_>>();
/// println!("{exp_prob:?}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/chisq/doctest_out/special_exp.md").unwrap();
/// # writeln!(f, "```output").unwrap();
/// # writeln!(f, "{chisq_dens:?}").unwrap();
/// # writeln!(f, "{exp_dens:?}").unwrap();
/// # writeln!(f, "{chisq_prob:?}").unwrap();
/// # writeln!(f, "{exp_prob:?}").unwrap();
/// # writeln!(f, "```").unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = include_str!("doctest_out/special_exp.md")))]
///
/// Visual testing; do P-P plots for 1000 points at various degrees of freedom
/// ```rust
/// # use r2rs_nmath::{
/// #     distribution::{ChiSquaredBuilder, NormalBuilder},
/// #     rng::MersenneTwister,
/// #     traits::{Distribution, RNG},
/// # };
/// # use r2rs_stats::funcs::ppoints;
/// # use strafe_plot::{
/// #     plot::Plot,
/// #     prelude::{IntoDrawingArea, Line, PlotOptions, Points, SVGBackend, BLACK, RED},
/// # };
/// # use strafe_type::FloatConstraint;
/// let l = 1.2;
/// let n = 1000;
/// let pp = ppoints(n, None)
///     .into_iter()
///     .map(|pp| pp.unwrap())
///     .collect::<Vec<_>>();
///
/// let mut rng = MersenneTwister::new();
/// rng.set_seed(1);
///
/// let norm = NormalBuilder::new().build();
/// let dfs = (0..9)
///     .map(|_| 2.0_f64.powf(4.0 * norm.random_sample(&mut rng).unwrap()))
///     .collect::<Vec<_>>();
///
/// let root = SVGBackend::new("df.svg", (1024, 768)).into_drawing_area();
/// for i in 0..9 {
///     let chisq = ChiSquaredBuilder::new().with_df(dfs[i]).with_ncp(l).build();
///     let mut rr = (0..n)
///         .map(|_| {
///             chisq
///                 .probability(chisq.random_sample(&mut rng).unwrap(), true)
///                 .unwrap()
///         })
///         .collect::<Vec<_>>();
///     rr.sort_by(|r1, r2| r1.partial_cmp(r2).unwrap());
///
///     Plot::new()
///         .with_options(PlotOptions {
///             title: format!("df={:.1}", dfs[i]),
///             x_axis_label: "pp".to_string(),
///             y_axis_label: "chisq.probability(chisq.random_sample())".to_string(),
///             plot_top: ((i / 3) % 3) as f64 / 3.0,
///             plot_bottom: (((i / 3) % 3) + 1) as f64 / 3.0,
///             plot_left: (i % 3) as f64 / 3.0,
///             plot_right: ((i % 3) + 1) as f64 / 3.0,
///             ..Default::default()
///         })
///         .with_plottable(Points {
///             x: pp.clone(),
///             y: rr,
///             color: BLACK,
///             ..Default::default()
///         })
///         .with_plottable(Line {
///             x: vec![-5.0, 5.0],
///             y: vec![-5.0, 5.0],
///             force_fit_all: false,
///             color: RED,
///             ..Default::default()
///         })
///         .plot(&root)
///         .unwrap();
/// }
/// # drop(root);
/// # use std::fs::rename;
/// # rename(
/// #     format!("df.svg"),
/// #     format!("src/distribution/chisq/doctest_out/df.svg"),
/// # )
/// # .unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = embed_doc_image::embed_image!("df", "src/distribution/chisq/doctest_out/df.svg")))]
#[cfg_attr(
    feature = "doc_outputs",
    cfg_attr(all(), doc = "![Degrees of Freedom Plot][df]")
)]
///
/// ```rust
/// # use num_traits::Float;
/// # use r2rs_base::traits::StatisticalSlice;
/// # use r2rs_nmath::{distribution::ChiSquaredBuilder, traits::Distribution};
/// # use strafe_type::FloatConstraint;
/// let lam = <[f64]>::sequence_by(0.0, 100.0, 0.25);
/// let exp_lam = lam.iter().map(|lam| (-lam / 2.0).exp()).collect::<Vec<_>>();
/// let p00 = lam
///     .iter()
///     .map(|lam| {
///         ChiSquaredBuilder::new()
///             .with_ncp(lam)
///             .with_df(f64::min_positive_value())
///             .build()
///             .probability(0.0, true)
///             .unwrap()
///     })
///     .collect::<Vec<_>>();
/// let p_0 = lam
///     .iter()
///     .map(|lam| {
///         ChiSquaredBuilder::new()
///             .with_ncp(lam)
///             .with_df(f64::min_positive_value())
///             .build()
///             .probability(1e-300, true)
///             .unwrap()
///     })
///     .collect::<Vec<_>>();
/// println!("{exp_lam:?}");
/// println!("{p00:?}");
/// println!("{p_0:?}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/chisq/doctest_out/analytical_test.md").unwrap();
/// # writeln!(f, "```output").unwrap();
/// # writeln!(f, "{exp_lam:?}").unwrap();
/// # writeln!(f, "{p00:?}").unwrap();
/// # writeln!(f, "{p_0:?}").unwrap();
/// # writeln!(f, "```").unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = include_str!("doctest_out/analytical_test.md")))]
pub struct ChiSquared {
    df: Rational64,
    ncp: Option<Positive64>,
}

impl Distribution for ChiSquared {
    fn density<R: Into<Real64>>(&self, x: R) -> Real64 {
        if let Some(ncp) = self.ncp {
            dnchisq(x, self.df, ncp, false)
        } else {
            dchisq(x, self.df, false)
        }
    }

    fn log_density<R: Into<Real64>>(&self, x: R) -> Real64 {
        if let Some(ncp) = self.ncp {
            dnchisq(x, self.df, ncp, true)
        } else {
            dchisq(x, self.df, true)
        }
    }

    fn probability<R: Into<Real64>>(&self, q: R, lower_tail: bool) -> Probability64 {
        if let Some(ncp) = self.ncp {
            pnchisq(q, self.df, ncp, lower_tail)
        } else {
            pchisq(q, self.df, lower_tail)
        }
    }

    fn log_probability<R: Into<Real64>>(&self, q: R, lower_tail: bool) -> LogProbability64 {
        if let Some(ncp) = self.ncp {
            log_pnchisq(q, self.df, ncp, lower_tail)
        } else {
            log_pchisq(q, self.df, lower_tail)
        }
    }

    fn quantile<P: Into<Probability64>>(&self, p: P, lower_tail: bool) -> Real64 {
        if let Some(ncp) = self.ncp {
            qnchisq(p, self.df, ncp, lower_tail)
        } else {
            qchisq(p, self.df, lower_tail)
        }
    }

    fn log_quantile<LP: Into<LogProbability64>>(&self, p: LP, lower_tail: bool) -> Real64 {
        if let Some(ncp) = self.ncp {
            log_qnchisq(p, self.df, ncp, lower_tail)
        } else {
            log_qchisq(p, self.df, lower_tail)
        }
    }

    fn random_sample<R: RNG>(&self, rng: &mut R) -> Real64 {
        if let Some(ncp) = self.ncp {
            rnchisq(self.df, ncp, rng)
        } else {
            rchisq(self.df, rng)
        }
    }
}

pub struct ChiSquaredBuilder {
    df: Option<Rational64>,
    ncp: Option<Positive64>,
}

impl ChiSquaredBuilder {
    pub fn new() -> Self {
        Self {
            df: None,
            ncp: None,
        }
    }

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

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

    pub fn build(&self) -> ChiSquared {
        let df = self.df.unwrap_or(1.0.into());

        ChiSquared { df, ncp: self.ncp }
    }
}

#[cfg(test)]
mod tests;

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

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