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 F Distribution
///
/// ## Description:
///
/// Density, distribution function, quantile function and random
/// generation for the F distribution with ‘df1’ and ‘df2’ degrees of
/// freedom (and optional non-centrality parameter ‘ncp’).
///
/// ## Arguments:
///
/// * df1, df2: degrees of freedom.  ‘Inf’ is allowed.
/// * ncp: non-centrality parameter. If omitted the central F is
/// assumed.
///
/// ## Details:
///
/// The F distribution with ‘df1 =’ n1 and ‘df2 =’ n2 degrees of
/// freedom has density
///
/// $f(x) = \frac{\Gamma(\frac{n1 + n2}{2})}{\Gamma(\frac{n1}{2}) \Gamma(\frac{n2}{2})}
/// (\frac{n1}{n2})^{\frac{n1}{2}} x^{\frac{n1}{2} - 1}
/// (1 + \frac{n1}{n2} x)^{-\frac{n1 + n2}{2}}$
///
/// for $x > 0$.
///
/// It is the distribution of the ratio of the mean squares of n1 and
/// n2 independent standard normals, and hence of the ratio of two
/// independent chi-squared variates each divided by its degrees of
/// freedom.  Since the ratio of a normal and the root mean-square of
/// m independent normals has a Student's t_m distribution, the square
/// of a t_m variate has a F distribution on 1 and m degrees of
/// freedom.
///
/// The non-central F distribution is again the ratio of mean squares
/// of independent normals of unit variance, but those in the
/// numerator are allowed to have non-zero means and ‘ncp’ is the sum
/// of squares of the means.  See Chisquare for further details on
/// non-central distributions.
///
/// ## Value:
///
/// ‘df’ gives the density, ‘pf’ gives the distribution function ‘qf’
/// gives the quantile function, and ‘rf’ generates random deviates.
///
/// Invalid arguments will result in return value ‘NaN’, with a
/// warning.
///
/// The length of the result is determined by ‘n’ for ‘rf’, and is the
/// maximum of the lengths of the numerical arguments for the other
/// functions.
///
/// The numerical arguments other than ‘n’ are recycled to the length
/// of the result.  Only the first elements of the logical arguments
/// are used.
///
/// ## Density Plot
///
/// ```rust
/// # use r2rs_base::traits::StatisticalSlice;
/// # use r2rs_nmath::{distribution::FBuilder, traits::Distribution};
/// # use strafe_plot::prelude::{IntoDrawingArea, Line, Plot, PlotOptions, SVGBackend, BLACK};
/// # use strafe_type::FloatConstraint;
/// let f = FBuilder::new().build();
/// let x = <[f64]>::sequence(-0.5, 4.0, 1000);
/// let y = x
///     .iter()
///     .map(|x| f.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/f/doctest_out/density.svg"),
/// #     )
/// #     .unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = embed_doc_image::embed_image!("density", "src/distribution/f/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:
///
/// For the central case of ‘df’, computed _via_ a binomial
/// probability, code contributed by Catherine Loader (see ‘dbinom’);
/// for the non-central case computed _via_ ‘dbeta’, code contributed
/// by Peter Ruckdeschel.
///
/// For ‘pf’, _via_ ‘pbeta’ (or for large ‘df2’, _via_ ‘pchisq’).
///
/// For ‘qf’, _via_ ‘qchisq’ for large ‘df2’, else _via_ ‘qbeta’.
///
/// ## 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_, volume 2, chapters 27 and 30.  Wiley,
/// New York.
///
/// ## See Also:
///
/// Distributions for other standard distributions, including ‘dchisq’
/// for chi-squared and ‘dt’ for Student's t distributions.
///
/// ## Examples:
///
/// Lower Tails
/// ```rust
/// # use r2rs_base::traits::StatisticalSlice;
/// # use r2rs_nmath::{
/// #     distribution::{FBuilder, TBuilder},
/// #     traits::Distribution,
/// # };
/// # use strafe_type::FloatConstraint;
/// let x = <[f64]>::sequence(0.001, 5.0, 100);
/// let nu = 4.0;
///
/// let t = TBuilder::new().with_df(nu).unwrap().build();
/// let f = FBuilder::new().with_df1(1).with_df2(nu).build();
/// let r1 = x
///     .iter()
///     .map(|x| 2.0 * t.probability(x, true).unwrap() - 1.0)
///     .collect::<Vec<_>>();
/// let r2 = x
///     .iter()
///     .map(|x| f.probability(x.powi(2), true).unwrap())
///     .collect::<Vec<_>>();
/// println!("{r1:?}");
/// println!("{r2:?}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/f/doctest_out/lower_tail.md").unwrap();
/// # writeln!(f, "```output").unwrap();
/// # writeln!(f, "{r1:?}").unwrap();
/// # writeln!(f, "{r2:?}").unwrap();
/// # writeln!(f, "```").unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = include_str!("doctest_out/lower_tail.md")))]
///
/// Upper Tails
/// ```rust
/// # use r2rs_base::traits::StatisticalSlice;
/// # use r2rs_nmath::{
/// #     distribution::{FBuilder, TBuilder},
/// #     traits::Distribution,
/// # };
/// # use strafe_type::FloatConstraint;
/// let x = <[f64]>::sequence(0.001, 5.0, 100);
/// let nu = 4.0;
///
/// let t = TBuilder::new().with_df(nu).unwrap().build();
/// let f = FBuilder::new().with_df1(1).with_df2(nu).build();
/// let r1 = x
///     .iter()
///     .map(|x| 2.0 * t.probability(x, false).unwrap())
///     .collect::<Vec<_>>();
/// let r2 = x
///     .iter()
///     .map(|x| f.probability(x.powi(2), false).unwrap())
///     .collect::<Vec<_>>();
/// println!("{r1:?}");
/// println!("{r2:?}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/f/doctest_out/upper_tail.md").unwrap();
/// # writeln!(f, "```output").unwrap();
/// # writeln!(f, "{r1:?}").unwrap();
/// # writeln!(f, "{r2:?}").unwrap();
/// # writeln!(f, "```").unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = include_str!("doctest_out/upper_tail.md")))]
///
/// The density of the square of a $t_m$ is $2*\frac{\{t-density}(x, m)}{2*x}$
///
/// Check this is the same as the density of $F_{1,m}$
/// ```rust
/// # use r2rs_base::traits::StatisticalSlice;
/// # use r2rs_nmath::{
/// #     distribution::{FBuilder, TBuilder},
/// #     traits::Distribution,
/// # };
/// # use strafe_type::FloatConstraint;
/// let x = <[f64]>::sequence(0.001, 5.0, 100);
/// let nu = 5.0;
///
/// let t = TBuilder::new().with_df(nu).unwrap().build();
/// let f = FBuilder::new().with_df1(1).with_df2(nu).build();
/// let r1 = x
///     .iter()
///     .map(|x| t.density(x).unwrap() / x)
///     .collect::<Vec<_>>();
/// let r2 = x
///     .iter()
///     .map(|x| f.density(x.powi(2)).unwrap())
///     .collect::<Vec<_>>();
/// println!("{r1:?}");
/// println!("{r2:?}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/f/doctest_out/dens.md").unwrap();
/// # writeln!(f, "```output").unwrap();
/// # writeln!(f, "{r1:?}").unwrap();
/// # writeln!(f, "{r2:?}").unwrap();
/// # writeln!(f, "```").unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = include_str!("doctest_out/dens.md")))]
///
/// Identity:  qf(2*p - 1, 1, df) == qt(p, df)^2  for  p >= 1/2
/// ```rust
/// use r2rs_base::traits::{QuantileType, StatisticalSlice};
/// use r2rs_nmath::{
///     distribution::{FBuilder, TBuilder},
///     traits::Distribution,
/// };
/// use strafe_type::FloatConstraint;
///
/// let p = <[f64]>::sequence(0.5, 0.99, 50);
/// let df = 10.0;
///
/// let rel_err = |x: &[f64], y: &[f64]| {
///     let mean = x
///         .iter()
///         .chain(y.iter())
///         .map(|f| f.abs())
///         .collect::<Vec<_>>()
///         .mean();
///     x.iter()
///         .zip(y.iter())
///         .map(|(x, y)| if x == y { 0.0 } else { (x - y).abs() / mean })
///         .collect::<Vec<_>>()
/// };
///
/// let t = TBuilder::new().with_df(df).unwrap().build();
/// let f = FBuilder::new().with_df1(1).with_df2(df).build();
/// let r1 = p
///     .iter()
///     .map(|p| t.quantile(p, true).unwrap().powi(2))
///     .collect::<Vec<_>>();
/// let r2 = p
///     .iter()
///     .map(|p| f.quantile(2.0 * p - 1.0, true).unwrap())
///     .collect::<Vec<_>>();
///
/// let error = rel_err(&r1, &r2).quantile(&[0.9], QuantileType::S);
///
/// println!("{error:?}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/f/doctest_out/quan.md").unwrap();
/// # writeln!(f, "```output").unwrap();
/// # writeln!(f, "{error:?}").unwrap();
/// # writeln!(f, "```").unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = include_str!("doctest_out/quan.md")))]
pub struct F {
    df1: Rational64,
    df2: Rational64,
    ncp: Option<Positive64>,
}

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

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

    fn probability<R: Into<Real64>>(&self, q: R, lower_tail: bool) -> Probability64 {
        if let Some(ncp) = self.ncp {
            pnf(q, self.df1, self.df2, ncp, lower_tail)
        } else {
            pf(q, self.df1, self.df2, lower_tail)
        }
    }

    fn log_probability<R: Into<Real64>>(&self, q: R, lower_tail: bool) -> LogProbability64 {
        if let Some(ncp) = self.ncp {
            log_pnf(q, self.df1, self.df2, ncp, lower_tail)
        } else {
            log_pf(q, self.df1, self.df2, lower_tail)
        }
    }

    fn quantile<P: Into<Probability64>>(&self, p: P, lower_tail: bool) -> Real64 {
        if let Some(ncp) = self.ncp {
            qnf(p, self.df1, self.df2, ncp, lower_tail)
        } else {
            qf(p, self.df1, self.df2, lower_tail)
        }
    }

    fn log_quantile<LP: Into<LogProbability64>>(&self, p: LP, lower_tail: bool) -> Real64 {
        if let Some(ncp) = self.ncp {
            log_qnf(p, self.df1, self.df2, ncp, lower_tail)
        } else {
            log_qf(p, self.df1, self.df2, lower_tail)
        }
    }

    fn random_sample<R: RNG>(&self, rng: &mut R) -> Real64 {
        if let Some(ncp) = self.ncp {
            rnf(self.df1, self.df2, ncp, rng)
        } else {
            rf(self.df1, self.df2, rng)
        }
    }
}

pub struct FBuilder {
    df1: Option<Rational64>,
    df2: Option<Rational64>,
    ncp: Option<Positive64>,
}

impl FBuilder {
    pub fn new() -> Self {
        Self {
            df1: None,
            df2: None,
            ncp: None,
        }
    }

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

    pub fn with_df2<R: Into<Rational64>>(&mut self, df2: R) -> &mut Self {
        self.df2 = Some(df2.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) -> F {
        let df1 = self.df1.unwrap_or(1.0.into());
        let df2 = self.df2.unwrap_or(1.0.into());

        F {
            df1,
            df2,
            ncp: self.ncp,
        }
    }
}

#[cfg(test)]
mod tests;

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

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