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

/// Implements the [Erlang](https://en.wikipedia.org/wiki/Erlang_distribution)
/// distribution
/// which is a special case of the
/// [Gamma](https://en.wikipedia.org/wiki/Gamma_distribution)
/// distribution
///
/// # Examples
///
/// ```
/// use stats::distribution::{Erlang, Continuous};
/// use stats::statistics::Distribution;
/// use math::assert_almost_eq;
///
/// let n = Erlang::new(3, 1.0).unwrap();
/// assert_eq!(n.mean().unwrap(), 3.0);
/// assert_almost_eq!(n.pdf(2.0), 0.2706705664732254, epsilon = 1e-15);
/// ```
#[derive(Copy, Clone, PartialEq, Debug)]
pub struct Erlang {
	g: Gamma,
}

impl Erlang {
	/// Constructs a new erlang distribution with a shape (k)
	/// of `shape` and a rate (λ) of `rate`
	///
	/// # Errors
	///
	/// Returns an error if `shape` or `rate` are `NaN`.
	/// Also returns an error if `shape == 0` or `rate <= 0.0`
	///
	/// # Examples
	///
	/// ```
	/// use stats::distribution::Erlang;
	///
	/// let mut result = Erlang::new(3, 1.0);
	/// assert!(result.is_ok());
	///
	/// result = Erlang::new(0, 0.0);
	/// assert!(result.is_err());
	/// ```
	pub fn new(shape: u64, rate: f64) -> Result<Erlang, GammaError> {
		Gamma::new(shape as f64, rate).map(|g| Erlang { g })
	}

	/// Returns the shape (k) of the erlang distribution
	///
	/// # Examples
	///
	/// ```
	/// use stats::distribution::Erlang;
	///
	/// let n = Erlang::new(3, 1.0).unwrap();
	/// assert_eq!(n.shape(), 3);
	/// ```
	pub fn shape(&self) -> u64 {
		self.g.shape() as u64
	}

	/// Returns the rate (λ) of the erlang distribution
	///
	/// # Examples
	///
	/// ```
	/// use stats::distribution::Erlang;
	///
	/// let n = Erlang::new(3, 1.0).unwrap();
	/// assert_eq!(n.rate(), 1.0);
	/// ```
	pub fn rate(&self) -> f64 {
		self.g.rate()
	}
}

impl core::fmt::Display for Erlang {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		write!(f, "E({}, {})", self.rate(), self.shape())
	}
}

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

impl ContinuousCDF for Erlang {
	/// Calculates the cumulative distribution function for the erlang
	/// distribution
	/// at `x`
	///
	/// # Formula
	///
	/// ```text
	/// γ(k, λx)  (k - 1)!
	/// ```
	///
	/// where `k` is the shape, `λ` is the rate, and `γ` is the lower
	/// incomplete gamma function
	fn cdf(&self, x: f64) -> f64 {
		self.g.cdf(x)
	}

	/// Calculates the cumulative distribution function for the erlang
	/// distribution
	/// at `x`
	///
	/// # Formula
	///
	/// ```text
	/// γ(k, λx)  (k - 1)!
	/// ```
	///
	/// where `k` is the shape, `λ` is the rate, 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 erlang
	/// distribution at `x`
	///
	/// # Formula
	///
	/// ```text
	/// γ^{-1}(k, (k - 1)! x) / λ
	/// ```
	///
	/// where `k` is the shape, `λ` is the rate, and `γ` is the upper
	/// incomplete gamma function
	fn inverse_cdf(&self, p: f64) -> f64 {
		self.g.inverse_cdf(p)
	}

	fn lower(&self) -> f64 {
		self.g.lower()
	}

	fn upper(&self) -> f64 {
		self.g.upper()
	}
}

impl Distribution for Erlang {
	/// Returns the mean of the erlang distribution
	///
	/// # Remarks
	///
	/// Returns `shape` if `rate == f64::INFINITY`. This behavior
	/// is borrowed from the Math.NET implementation
	///
	/// # Formula
	///
	/// ```text
	/// k / λ
	/// ```
	///
	/// where `k` is the shape and `λ` is the rate
	fn mean(&self) -> Option<f64> {
		self.g.mean()
	}

	/// Returns the variance of the erlang distribution
	///
	/// # Formula
	///
	/// ```text
	/// k / λ^2
	/// ```
	///
	/// where `α` is the shape and `λ` is the rate
	fn variance(&self) -> Option<f64> {
		self.g.variance()
	}

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

	/// Returns the skewness of the erlang distribution
	///
	/// # Formula
	///
	/// ```text
	/// 2 / sqrt(k)
	/// ```
	///
	/// where `k` is the shape
	fn skewness(&self) -> Option<f64> {
		self.g.skewness()
	}
}

impl Mode<Option<f64>> for Erlang {
	/// Returns the mode for the erlang distribution
	///
	/// # Remarks
	///
	/// Returns `shape` if `rate ==f64::INFINITY`. This behavior
	/// is borrowed from the Math.NET implementation
	///
	/// # Formula
	///
	/// ```text
	/// (k - 1) / λ
	/// ```
	///
	/// where `k` is the shape and `λ` is the rate
	fn mode(&self) -> Option<f64> {
		self.g.mode()
	}
}

impl Continuous for Erlang {
	type T = f64;

	/// Calculates the probability density function for the erlang distribution
	/// at `x`
	///
	/// # Remarks
	///
	/// Returns `NAN` if any of `shape` or `rate` are `INF`
	/// or if `x` is `INF`
	///
	/// # Formula
	///
	/// ```text
	/// (λ^k / Γ(k)) * x^(k - 1) * e^(-λ * x)
	/// ```
	///
	/// where `k` is the shape, `λ` is the rate, and `Γ` is the gamma function
	fn pdf(&self, x: f64) -> f64 {
		self.g.pdf(x)
	}

	/// Calculates the log probability density function for the erlang
	/// distribution
	/// at `x`
	///
	/// # Remarks
	///
	/// Returns `NAN` if any of `shape` or `rate` are `INF`
	/// or if `x` is `INF`
	///
	/// # Formula
	///
	/// ```text
	/// ln((λ^k / Γ(k)) * x^(k - 1) * e ^(-λ * x))
	/// ```
	///
	/// where `k` is the shape, `λ` is the rate, and `Γ` is the gamma function
	fn ln_pdf(&self, x: f64) -> f64 {
		self.g.ln_pdf(x)
	}
}
