use std::fmt;
use std::ops;

use crate::{utils, MU, SIGMA};

#[derive(Clone, Copy, PartialEq, Debug)]
pub struct Gaussian {
    mu: f64,
    sigma: f64,
}

impl Gaussian {
    pub const fn new(mu: f64, sigma: f64) -> Self {
        Gaussian { mu, sigma }
    }

    pub fn mu(&self) -> f64 {
        self.mu
    }

    pub fn sigma(&self) -> f64 {
        self.sigma
    }

    pub fn tau(&self) -> f64 {
        self.mu * self.pi()
    }

    pub fn pi(&self) -> f64 {
        self.sigma.powi(-2)
    }

    pub fn forget(&self, gamma: f64, t: f64) -> Self {
        Self::new(self.mu, (self.sigma().powi(2) + t * gamma.powi(2)).sqrt())
    }

    pub fn delta(&self, m: Gaussian) -> (f64, f64) {
        ((self.mu() - m.mu()).abs(), (self.sigma() - m.sigma()).abs())
    }

    pub fn exclude(&self, m: Gaussian) -> Gaussian {
        Gaussian::new(
            self.mu() - m.mu(),
            (self.sigma().powi(2) - m.sigma().powi(2)).sqrt(),
        )
    }

    /*
    def forget(self,gamma,t):
        return Gaussian(self.mu, math.sqrt(self.sigma**2 + t*gamma**2))

    def delta(self, M):
        return abs(self.mu - M.mu) , abs(self.sigma - M.sigma)

    def exclude(self, M):
        return Gaussian(self.mu - M.mu, math.sqrt(self.sigma**2 - M.sigma**2) )

    def isapprox(self, M, tol=1e-4):
        return (abs(self.mu - M.mu) < tol) and (abs(self.sigma - M.sigma) < tol)
    */
}

impl Default for Gaussian {
    fn default() -> Self {
        Gaussian {
            mu: MU,
            sigma: SIGMA,
        }
    }
}

impl fmt::Display for Gaussian {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "N(mu={:.3}, sigma={:.3})", self.mu, self.sigma)
    }
}

impl ops::Add<Gaussian> for Gaussian {
    type Output = Gaussian;

    fn add(self, rhs: Gaussian) -> Self::Output {
        Gaussian {
            mu: self.mu + rhs.mu,
            sigma: (self.sigma.powi(2) + rhs.sigma.powi(2)).sqrt(),
        }
    }
}

impl ops::Sub<Gaussian> for Gaussian {
    type Output = Gaussian;

    fn sub(self, rhs: Gaussian) -> Self::Output {
        Gaussian {
            mu: self.mu - rhs.mu,
            sigma: (self.sigma.powi(2) + rhs.sigma.powi(2)).sqrt(),
        }
    }
}

impl ops::Mul<Gaussian> for Gaussian {
    type Output = Gaussian;

    fn mul(self, rhs: Gaussian) -> Self::Output {
        let (mu, sigma) = utils::mu_sigma(self.tau() + rhs.tau(), self.pi() + rhs.pi());

        Gaussian { mu, sigma }
    }
}

impl ops::Div<Gaussian> for Gaussian {
    type Output = Gaussian;

    fn div(self, rhs: Gaussian) -> Self::Output {
        let (mu, sigma) = utils::mu_sigma(self.tau() - rhs.tau(), self.pi() - rhs.pi());

        Gaussian { mu, sigma }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_add() {
        let n = Gaussian {
            mu: 25.0,
            sigma: 25.0 / 3.0,
        };

        let m = Gaussian {
            mu: 0.0,
            sigma: 1.0,
        };

        assert_eq!(
            n + m,
            Gaussian {
                mu: 25.0,
                sigma: 8.393118874676116
            }
        );
    }

    #[test]
    fn test_sub() {
        let n = Gaussian {
            mu: 25.0,
            sigma: 25.0 / 3.0,
        };

        let m = Gaussian {
            mu: 1.0,
            sigma: 1.0,
        };

        assert_eq!(
            n - m,
            Gaussian {
                mu: 24.0,
                sigma: 8.393118874676116
            }
        );
    }

    #[test]
    fn test_mul() {
        let n = Gaussian {
            mu: 25.0,
            sigma: 25.0 / 3.0,
        };

        let m = Gaussian {
            mu: 0.0,
            sigma: 1.0,
        };

        assert_eq!(
            n * m,
            Gaussian {
                mu: 0.35488958990536273,
                sigma: 0.992876838486922
            }
        );
    }

    #[test]
    fn test_div() {
        let n = Gaussian {
            mu: 25.0,
            sigma: 25.0 / 3.0,
        };

        let m = Gaussian {
            mu: 0.0,
            sigma: 1.0,
        };

        assert_eq!(
            m / n,
            Gaussian {
                mu: -0.3652597402597402,
                sigma: 1.0072787050317253
            }
        );
    }
}