use crate::distribution::{Continuous, ContinuousCDF, Gamma, GammaError};
use crate::statistics::*;
use core::f64;

/// Implements the
/// [Chi-squared](https://en.wikipedia.org/wiki/Chi-squared_distribution)
/// distribution which is a special case of the
/// [Gamma](https://en.wikipedia.org/wiki/Gamma_distribution) distribution
/// (referenced [Here](./struct.Gamma.html))
///
/// # Examples
///
/// ```
/// use stats::distribution::{ChiSquared, Continuous};
/// use stats::statistics::Distribution;
/// use math::assert_almost_eq;
///
/// let n = ChiSquared::new(3.0).unwrap();
/// assert_eq!(n.mean().unwrap(), 3.0);
/// assert_almost_eq!(n.pdf(4.0), 0.1079819330263761, epsilon = 1e-15);
/// ```
#[derive(Copy, Clone, PartialEq, Debug)]
pub struct ChiSquared {
	freedom: f64,
	g: Gamma,
}

impl ChiSquared {
	/// Constructs a new chi-squared distribution with `freedom`
	/// degrees of freedom. This is equivalent to a Gamma distribution
	/// with a shape of `freedom / 2.0` and a rate of `0.5`.
	///
	/// # Errors
	///
	/// Returns an error if `freedom` is `NaN` or less than
	/// or equal to `0.0`
	///
	/// # Examples
	///
	/// ```
	/// use stats::distribution::ChiSquared;
	///
	/// let mut result = ChiSquared::new(3.0);
	/// assert!(result.is_ok());
	///
	/// result = ChiSquared::new(0.0);
	/// assert!(result.is_err());
	/// ```
	pub fn new(freedom: f64) -> Result<ChiSquared, GammaError> {
		Gamma::new(freedom / 2.0, 0.5)
			.map(|g| ChiSquared { freedom, g })
	}

	/// Returns the degrees of freedom of the chi-squared
	/// distribution
	///
	/// # Examples
	///
	/// ```
	/// use stats::distribution::ChiSquared;
	///
	/// let n = ChiSquared::new(3.0).unwrap();
	/// assert_eq!(n.freedom(), 3.0);
	/// ```
	pub fn freedom(&self) -> f64 {
		self.freedom
	}

	/// Returns the shape of the underlying Gamma distribution
	///
	/// # Examples
	///
	/// ```
	/// use stats::distribution::ChiSquared;
	///
	/// let n = ChiSquared::new(3.0).unwrap();
	/// assert_eq!(n.shape(), 3.0 / 2.0);
	/// ```
	pub fn shape(&self) -> f64 {
		self.g.shape()
	}

	/// Returns the rate of the underlying Gamma distribution
	///
	/// # Examples
	///
	/// ```
	/// use stats::distribution::ChiSquared;
	///
	/// let n = ChiSquared::new(3.0).unwrap();
	/// assert_eq!(n.rate(), 0.5);
	/// ```
	pub fn rate(&self) -> f64 {
		self.g.rate()
	}
}

impl core::fmt::Display for ChiSquared {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		write!(f, "χ^2_{}", self.freedom)
	}
}

#[cfg(feature = "rand")]
#[cfg_attr(docsrs, doc(cfg(feature = "rand")))]
impl rand::distr::Distribution<f64> for ChiSquared {
	fn sample<R: rand::Rng + ?Sized>(&self, r: &mut R) -> f64 {
		rand::distr::Distribution::sample(&self.g, r)
	}
}

impl ContinuousCDF for ChiSquared {
	/// Calculates the cumulative distribution function for the
	/// chi-squared distribution at `x`
	///
	/// # Formula
	///
	/// ```text
	/// (1 / Γ(k / 2)) * γ(k / 2, x / 2)
	/// ```
	///
	/// where `k` is the degrees of freedom, `Γ` is the gamma function,
	/// and `γ` is the lower incomplete gamma function
	fn cdf(&self, x: f64) -> f64 {
		self.g.cdf(x)
	}

	/// Calculates the cumulative distribution function for the
	/// chi-squared distribution at `x`
	///
	/// # Formula
	///
	/// ```text
	/// (1 / Γ(k / 2)) * γ(k / 2, x / 2)
	/// ```
	///
	/// where `k` is the degrees of freedom, `Γ` is the gamma function,
	/// and `γ` is the upper incomplete gamma function
	fn sf(&self, x: f64) -> f64 {
		self.g.sf(x)
	}

	/// Calculates the inverse cumulative distribution function for the
	/// chi-squared distribution at `x`
	///
	/// # Formula
	///
	/// ```text
	/// γ^{-1}(k / 2, x * Γ(k / 2) / 2)
	/// ```
	///
	/// where `k` is the degrees of freedom, `Γ` is the gamma function,
	/// and `γ` is the lower incomplete gamma function
	fn inverse_cdf(&self, p: f64) -> f64 {
		self.g.inverse_cdf(p)
	}

	fn lower(&self) -> f64 {
		0.0
	}

	fn upper(&self) -> f64 {
		f64::INFINITY
	}
}

impl Distribution for ChiSquared {
	/// Returns the mean of the chi-squared distribution
	///
	/// # Formula
	///
	/// ```text
	/// k
	/// ```
	///
	/// where `k` is the degrees of freedom
	fn mean(&self) -> Option<f64> {
		self.g.mean()
	}

	/// Returns the median  of the chi-squared distribution
	///
	/// # Formula
	///
	/// ```text
	/// k * (1 - (2 / 9k))^3
	/// ```
	fn median(&self) -> Option<f64> {
		let median = if self.freedom < 1.0 {
			// if k is small, calculate using expansion of formula
			self.freedom - 2.0 / 3.0 + 12.0 / (81.0 * self.freedom)
				- 8.0 / (729.0 * self.freedom * self.freedom)
		} else {
			// if k is large enough, median heads toward k - 2/3
			self.freedom - 2.0 / 3.0
		};
		Some(median)
	}

	/// Returns the variance of the chi-squared distribution
	///
	/// # Formula
	///
	/// ```text
	/// 2k
	/// ```
	///
	/// where `k` is the degrees of freedom
	fn variance(&self) -> Option<f64> {
		self.g.variance()
	}

	/// Returns the entropy of the chi-squared distribution
	///
	/// # Formula
	///
	/// ```text
	/// (k / 2) + ln(2 * Γ(k / 2)) + (1 - (k / 2)) * ψ(k / 2)
	/// ```
	///
	/// where `k` is the degrees of freedom, `Γ` is the gamma function,
	/// and `ψ` is the digamma function
	fn entropy(&self) -> Option<f64> {
		self.g.entropy()
	}

	/// Returns the skewness of the chi-squared distribution
	///
	/// # Formula
	///
	/// ```text
	/// sqrt(8 / k)
	/// ```
	///
	/// where `k` is the degrees of freedom
	fn skewness(&self) -> Option<f64> {
		self.g.skewness()
	}
}

impl Mode<Option<f64>> for ChiSquared {
	/// Returns the mode of the chi-squared distribution
	///
	/// # Formula
	///
	/// ```text
	/// k - 2
	/// ```
	///
	/// where `k` is the degrees of freedom
	fn mode(&self) -> Option<f64> {
		self.g.mode()
	}
}

impl Continuous for ChiSquared {
	type T = f64;

	/// Calculates the probability density function for the chi-squared
	/// distribution at `x`
	///
	/// # Formula
	///
	/// ```text
	/// 1 / (2^(k / 2) * Γ(k / 2)) * x^((k / 2) - 1) * e^(-x / 2)
	/// ```
	///
	/// where `k` is the degrees of freedom and `Γ` is the gamma function
	fn pdf(&self, x: f64) -> f64 {
		self.g.pdf(x)
	}

	/// Calculates the log probability density function for the chi-squared
	/// distribution at `x`
	///
	/// # Formula
	///
	/// ```text
	/// ln(1 / (2^(k / 2) * Γ(k / 2)) * x^((k / 2) - 1) * e^(-x / 2))
	/// ```
	fn ln_pdf(&self, x: f64) -> f64 {
		self.g.ln_pdf(x)
	}
}
