use ndarray::prelude::*;

use super::Kernel;

#[derive(Clone, Copy)]
pub struct Exponential {
    var: f64,
    l_scale: f64,
}

impl Exponential {
    pub fn new(var: f64, l_scale: f64) -> Self {
        Exponential { var, l_scale }
    }
}

impl Kernel for Exponential {
    fn k_mat(&self, ts1: &[f64], ts2: Option<&[f64]>) -> Array2<f64> {
        let ts2 = ts2.unwrap_or(ts1);

        let mut r = super::distance(ts1, ts2) / self.l_scale;
        r.mapv_inplace(|v| (-v).exp());

        r * self.var
    }

    fn k_diag(&self, ts: &[f64]) -> Array1<f64> {
        Array1::ones(ts.len()) * self.var
    }

    fn order(&self) -> usize {
        1
    }

    fn state_mean(&self, _t: f64) -> Array1<f64> {
        Array1::zeros(1)
    }

    fn state_cov(&self, _t: f64) -> Array2<f64> {
        array![[1.0]] * self.var
    }

    fn measurement_vector(&self) -> Array1<f64> {
        array![1.0]
    }

    fn feedback(&self) -> Array2<f64> {
        (-1.0 / self.l_scale) * array![[1.0]]
    }

    fn noise_effect(&self) -> Array2<f64> {
        array![[1.0]]
    }

    fn transition(&self, t0: f64, t1: f64) -> Array2<f64> {
        (-(t1 - t0) / self.l_scale).exp() * array![[1.0]]
    }

    fn noise_cov(&self, t0: f64, t1: f64) -> Array2<f64> {
        self.var * (1.0 - (-2.0 * (t1 - t0) / self.l_scale).exp()) * array![[1.0]]
    }
}

#[cfg(test)]
mod tests {
    extern crate blas_src;

    use approx::assert_abs_diff_eq;
    use rand::{distributions::Standard, thread_rng, Rng};

    use super::*;

    #[test]
    fn test_kernel_matrix() {
        let kernel = Exponential::new(1.1, 2.2);

        let ts = [1.26, 1.46, 2.67];

        assert_abs_diff_eq!(
            kernel.k_mat(&ts, None),
            array![
                [1.1, 1.0044107879104887, 0.5794946136290716],
                [1.0044107879104887, 1.1, 0.6346447914185355],
                [0.5794946136290716, 0.6346447914185355, 1.1]
            ]
        );
    }

    #[test]
    fn test_kernel_diag() {
        let kernel = Exponential::new(1.1, 2.2);

        let ts: Vec<_> = thread_rng()
            .sample_iter::<f64, _>(Standard)
            .take(10)
            .map(|x| x * 10.0)
            .collect();

        assert_eq!(kernel.k_mat(&ts, None).diag(), kernel.k_diag(&ts));
    }

    #[test]
    fn test_kernel_order() {
        let kernel = Exponential::new(1.1, 2.2);

        let m = kernel.order();

        assert_eq!(kernel.state_mean(0.0).shape(), &[m]);
        assert_eq!(kernel.state_cov(0.0).shape(), &[m, m]);
        assert_eq!(kernel.measurement_vector().shape(), &[m]);
        assert_eq!(kernel.feedback().shape(), &[m, m]);
        assert_eq!(kernel.noise_effect().shape()[0], m);
        assert_eq!(kernel.transition(0.0, 1.0).shape(), &[m, m]);
        assert_eq!(kernel.noise_cov(0.0, 1.0).shape(), &[m, m]);
    }

    #[test]
    fn test_ssm_variance() {
        let kernel = Exponential::new(1.1, 2.2);

        let ts: Vec<_> = thread_rng()
            .sample_iter::<f64, _>(Standard)
            .take(10)
            .map(|x| x * 10.0)
            .collect();

        let h = kernel.measurement_vector();

        let vars = ts
            .iter()
            .map(|t| h.dot(&kernel.state_cov(*t)).dot(&h))
            .collect::<Vec<_>>();

        assert_abs_diff_eq!(Array::from(vars), kernel.k_diag(&ts));
    }
}