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

use std::{
    error::Error,
    fmt::{Display, Error as FmtError, Formatter},
};

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

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

/// # The Student t Distribution
///
/// ## Description
///
/// Density, distribution function, quantile function and random generation for the t distribution
/// with df degrees of freedom (and optional non-centrality parameter ncp).
///
/// ## Arguments
///
/// * df: degrees of freedom (> 0, maybe non-integer). df = Inf is allowed.
/// * ncp: non-centrality parameter delta; currently except for rt(), only for abs(ncp) <= 37.62. If omitted, use the central t distribution.
///
/// ## Details
///
/// The t distribution with df = n degrees of freedom has density
///
/// $f(x) = \Gamma(\frac{n+1}{2}) / (\sqrt{n \pi} \Gamma(\frac{n}{2})) (1 + \frac{x^2}{n})^{-\frac{n+1}{2}}$
///
/// for all real x. It has mean 0 (for n > 1) and variance $\frac{n}{n-2}$ (for n > 2).
///
/// The general non-central t with parameters $(df, Del) = (df, ncp)$ is defined as the distribution
/// of $T(df, Del) := (U + Del) / \sqrt{V/df}$ where U and V are independent random variables,
/// $U$ ~ $N(0,1)$ and $V$ ~ $\chi^2(df)$ (see Chisquare).
///
/// The most used applications are power calculations for t-tests:
/// Let $T= \frac{mX - m0}{S/\sqrt{n}}$ where mX is the mean and S the sample standard deviation (sd)
/// of $X_1, X_2, …, X_n$ which are i.i.d. $N(\mu, \sigma^2)$ Then T is distributed as non-central t with
/// $df = n - 1$ degrees of freedom and non-centrality parameter $ncp = (\mu - m0) * \sqrt{n}/\sigma$.
///
/// ## Density Plot
///
/// ```rust
/// # use r2rs_base::traits::StatisticalSlice;
/// # use r2rs_nmath::{distribution::TBuilder, traits::Distribution};
/// # use strafe_plot::prelude::{IntoDrawingArea, Line, Plot, PlotOptions, SVGBackend, BLACK};
/// # use strafe_type::FloatConstraint;
/// let t = TBuilder::new().build();
/// let x = <[f64]>::sequence(-10.0, 10.0, 1000);
/// let y = x
///     .iter()
///     .map(|x| t.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/t/doctest_out/density.svg"),
/// #     )
/// #     .unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = embed_doc_image::embed_image!("density", "src/distribution/t/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 dt is computed via an accurate formula provided by Catherine Loader (see the
/// reference in dbinom).
///
/// For the non-central case of dt, C code contributed by Claus Ekstrøm based on the relationship
/// (for x != 0) to the cumulative distribution.
///
/// For the central case of pt, a normal approximation in the tails, otherwise via pbeta.
///
/// For the non-central case of pt based on a C translation of
///
/// Lenth, R. V. (1989). Algorithm AS 243 — Cumulative distribution function of the non-central t
/// distribution, Applied Statistics 38, 185–189.
///
/// This 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.
///
/// For central qt, a C translation of
///
/// Hill, G. W. (1970) Algorithm 396: Student's t-quantiles. Communications of the ACM, 13(10),
/// 619–620.
///
/// altered to take account of
///
/// Hill, G. W. (1981) Remark on Algorithm 396, ACM Transactions on Mathematical Software, 7, 250–1.
///
/// The non-central case is done by inversion.
///
/// ## References
///
/// Becker, R. A., Chambers, J. M. and Wilks, A. R. (1988) The New S Language. Wadsworth &
/// Brooks/Cole. (Except non-central versions.)
///
/// Johnson, N. L., Kotz, S. and Balakrishnan, N. (1995) Continuous Univariate Distributions,
/// volume 2, chapters 28 and 31. Wiley, New York.
///
/// ## See Also
///
/// Distributions for other standard distributions, including df for the F distribution.
///
/// ## Examples
///
/// ```rust
/// # use r2rs_nmath::{distribution::TBuilder, traits::Distribution};
/// # use strafe_type::FloatConstraint;
/// let x = (1..=5).collect::<Vec<_>>();
/// let t = TBuilder::new().with_df(1).unwrap().build();
/// let r = x
///     .iter()
///     .map(|x| 1.0 - t.probability(x, true).unwrap())
///     .collect::<Vec<_>>();
/// println!("{r:?}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/t/doctest_out/dens1.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/dens1.md")))]
///
/// ```rust
/// # use r2rs_nmath::{distribution::TBuilder, traits::Distribution};
/// # use strafe_type::FloatConstraint;
/// let dfs = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 20, 50, 100, 1000];
/// let r = dfs
///     .iter()
///     .map(|df| {
///         TBuilder::new()
///             .with_df(df)
///             .unwrap()
///             .build()
///             .quantile(0.975, true)
///             .unwrap()
///     })
///     .collect::<Vec<_>>();
/// println!("{r:?}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/t/doctest_out/dens2.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/dens2.md")))]
///
/// ```rust
/// # use r2rs_base::traits::StatisticalSlice;
/// # use r2rs_nmath::{distribution::TBuilder, traits::Distribution};
/// # use strafe_plot::prelude::{IntoDrawingArea, Line, Plot, PlotOptions, SVGBackend, BLACK};
/// # use strafe_type::FloatConstraint;
/// let nt = TBuilder::new().with_df(3).unwrap().with_ncp(2).build();
/// let x = <[f64]>::sequence_by(-3.0, 11.0, 0.01);
/// let y = x.iter().map(|x| nt.density(x).unwrap()).collect::<Vec<_>>();
///
/// let root = SVGBackend::new("noncent_density.svg", (1024, 768)).into_drawing_area();
/// Plot::new()
///     .with_options(PlotOptions {
///         x_axis_label: "x".to_string(),
///         y_axis_label: "density".to_string(),
///         title: "Non-Central T 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!("noncent_density.svg"),
/// #             format!("src/distribution/t/doctest_out/noncent_density.svg"),
/// #     )
/// #     .unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = embed_doc_image::embed_image!("noncent_density", "src/distribution/t/doctest_out/noncent_density.svg")))]
#[cfg_attr(
    feature = "doc_outputs",
    cfg_attr(all(), doc = "![Noncentral Density][noncent_density]")
)]
///
/// ```r
/// require(graphics)
///
/// tt <- seq(0, 10, len = 21)
/// ncp <- seq(0, 6, len = 31)
/// ptn <- outer(tt, ncp, function(t, d) pt(t, df = 3, ncp = d))
/// t.tit <- "Non-central t - Probabilities"
/// image(tt, ncp, ptn, zlim = c(0,1), main = t.tit)
/// persp(tt, ncp, ptn, zlim = 0:1, r = 2, phi = 20, theta = 200, main = t.tit,
///       xlab = "t", ylab = "non-centrality parameter",
///       zlab = "Pr(T <= t)")
/// ```
pub struct T {
    df: Rational64,
    ncp: Option<Real64>,
}

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

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

    fn probability<R: Into<Real64>>(&self, q: R, lower_tail: bool) -> Probability64 {
        if let Some(ncp) = self.ncp {
            pnt(q, self.df, ncp, lower_tail)
        } else {
            pt(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_pnt(q, self.df, ncp, lower_tail)
        } else {
            log_pt(q, self.df, lower_tail)
        }
    }

    fn quantile<P: Into<Probability64>>(&self, p: P, lower_tail: bool) -> Real64 {
        if let Some(ncp) = self.ncp {
            qnt(p, self.df, ncp, lower_tail)
        } else {
            qt(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_qnt(p, self.df, ncp, lower_tail)
        } else {
            log_qt(p, self.df, lower_tail)
        }
    }

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

#[derive(Debug)]
pub enum BuildError {
    DFLessThan0,
    BadFloat,
}

#[cfg(not(tarpaulin_include))]
impl Display for BuildError {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), FmtError> {
        match &self {
            BuildError::DFLessThan0 => write!(fmt, "DF must be greater than or equal to 0"),
            BuildError::BadFloat => write!(fmt, "The float passed was not valid"),
        }
    }
}

impl Error for BuildError {}

pub struct TBuilder {
    df: Option<Rational64>,
    ncp: Option<Real64>,
}

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

    pub fn with_df<R: Into<Rational64>>(&mut self, df: R) -> Result<&mut Self, impl Error> {
        let df = df.into();
        if df.unwrap() > 0.0 {
            self.df = Some(df);
            Ok(self)
        } else {
            #[cfg(not(tarpaulin_include))]
            Err(BuildError::DFLessThan0)
        }
    }

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

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

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

#[cfg(test)]
mod tests;

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

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