#![allow(non_snake_case)]

//! Provide FatEstimator trait to allow for testing of many estimators with disparate requirements.
//!
//! This defines predict and observe operations using linearised models to be tested.
//!
//! Where necessary 'fat' estimator states are defined so that different operations of an estimator can be tested using
//! the FatEstimator trait.
//!
//! Implementation are provides for all the linearised estimators.
//! The [`sir`] sampled estimator is implemented, so it can be tested using linearised models.

use std::iter::StepBy;

use nalgebra::{allocator::Allocator, DefaultAllocator, storage::Storage};
use nalgebra::{Dim, DimAdd, DimMin, DimMinimum, DimSum, U1};
use nalgebra::{Matrix, Scalar};
use nalgebra::{OMatrix, OVector, RealField};
use nalgebra::iter::MatrixIter;
use num_traits::FromPrimitive;
use rand::RngCore;

use bayes_estimate::estimators::information_root::InformationRootState;
use bayes_estimate::estimators::sir;
use bayes_estimate::estimators::sir::{SampleState, SampleStateDraw};
use bayes_estimate::estimators::ud::{CorrelatedFactorNoise, UDState};
use bayes_estimate::estimators::unscented_duplex::UnscentedDuplexState;
use bayes_estimate::models::{
    Estimator, ExtendedLinearObserver, ExtendedLinearPredictor, FunctionalObserver,
    FunctionalPredictor, InformationState, KalmanEstimator, KalmanState,
};
use bayes_estimate::noise::{CorrelatedNoise, CoupledNoise, UncorrelatedNoise};

/// Define the estimator operations to be tested.
pub trait FatEstimator<D: Dim, QD: Dim, ZD: Dim>:
    Estimator<f64, D> + KalmanEstimator<f64, D>
where
    DefaultAllocator: Allocator<D, D>
        + Allocator<D>
        + Allocator<QD, QD>
        + Allocator<D, QD>
        + Allocator<QD>
        + Allocator<ZD, ZD>
        + Allocator<ZD, D>
        + Allocator<ZD>,
{
    fn dim(&self) -> D {
        return Estimator::state(self).unwrap().shape_generic().0;
    }

    fn allow_error_by(&self) -> f64 {
        1f64
    }

    fn trace_state(&self) {}

    // Initialize with KalmanState
    fn init(&mut self, state: &KalmanState<f64, D>) -> Result<(), &str>;

    /// Prediction with additive noise
    fn predict_fn(
        &mut self,
        x_pred: &OVector<f64, D>,
        f: fn(&OVector<f64, D>) -> OVector<f64, D>,
        F: &OMatrix<f64, D, D>,
        noise: &CoupledNoise<f64, D, QD>,
    ) where
        DefaultAllocator: Allocator<D, U1> + Allocator<U1>;

    /// Observation with correlated noise
    fn observe_any(
        &mut self,
        z: &OVector<f64, ZD>,
        h: fn(&OVector<f64, D>) -> OVector<f64, ZD>,
        h_normalize: fn(h: &mut OVector<f64, ZD>, h0: &OVector<f64, ZD>),
        H: &OMatrix<f64, ZD, D>,
        noise: &CorrelatedNoise<f64, ZD>,
    ) -> Result<(), &str>;
}

/// Test covariance estimator operations defined on a KalmanState.
impl<D: Dim, QD: Dim, ZD: Dim> FatEstimator<D, QD, ZD> for KalmanState<f64, D>
where
    DefaultAllocator: Allocator<D, D>
        + Allocator<D>
        + Allocator<QD, QD>
        + Allocator<D, QD>
        + Allocator<QD>
        + Allocator<ZD, ZD>
        + Allocator<ZD, D>
        + Allocator<D, ZD>
        + Allocator<ZD>
        + Allocator<U1, U1>
        + Allocator<U1, D>
        + Allocator<U1>,
{
    fn init(&mut self, state: &KalmanState<f64, D>) -> Result<(), &str> {
        self.x = state.x.clone();
        self.X = state.X.clone();
        Ok(())
    }

    fn predict_fn(
        &mut self,
        x_pred: &OVector<f64, D>,
        _f: fn(&OVector<f64, D>) -> OVector<f64, D>,
        F: &OMatrix<f64, D, D>,
        noise: &CoupledNoise<f64, D, QD>,
    ) {
        ExtendedLinearPredictor::<f64, D>::predict(
            self,
            x_pred,
            F,
            &CorrelatedNoise::from_coupled::<QD>(noise),
        )
        .unwrap();
    }

    fn observe_any(
        &mut self,
        z: &OVector<f64, ZD>,
        h: fn(&OVector<f64, D>) -> OVector<f64, ZD>,
        h_normalize: fn(h: &mut OVector<f64, ZD>, h0: &OVector<f64, ZD>),
        H: &OMatrix<f64, ZD, D>,
        noise: &CorrelatedNoise<f64, ZD>,
    ) -> Result<(), &str> {
        let mut hx = h(&self.x);
        h_normalize(&mut hx, &z);
        let s = z - &hx;

        ExtendedLinearObserver::observe_innovation(self, &s, H, noise)
        // TODO tests application of observations sequentially
        // for o in 0..(s.nrows()) {
        //     ExtendedLinearObserver::observe_innovation(self, &Vector1::new(s[o]), &H.row(o).into_owned(), &CorrelatedNoise{Q: Matrix1::new(noise.Q[(o, o)])}).unwrap();
        // }
        // Ok(())
    }
}

/// Test information estimator operations defined on a InformationState.
impl<D: Dim, QD: Dim, ZD: Dim> FatEstimator<D, QD, ZD> for InformationState<f64, D>
where
    DefaultAllocator: Allocator<D, D>
        + Allocator<D>
        + Allocator<QD, QD>
        + Allocator<D, QD>
        + Allocator<QD>
        + Allocator<ZD, ZD>
        + Allocator<ZD, D>
        + Allocator<ZD>,
    // observe_innovation
    DefaultAllocator: Allocator<D, ZD>,
{
    fn dim(&self) -> D {
        return self.i.shape_generic().0;
    }

    fn init(&mut self, state: &KalmanState<f64, D>) -> Result<(), &str> {
        let info = InformationState::try_from(state.clone()).unwrap();
        self.i = info.i;
        self.I = info.I;
        Ok(())
    }

    fn predict_fn(
        &mut self,
        x_pred: &OVector<f64, D>,
        _f: fn(&OVector<f64, D>) -> OVector<f64, D>,
        F: &OMatrix<f64, D, D>,
        noise: &CoupledNoise<f64, D, QD>,
    ) {
        ExtendedLinearPredictor::<f64, D>::predict(
            self,
            x_pred,
            F,
            &CorrelatedNoise::from_coupled::<QD>(noise),
        )
        .unwrap();
    }

    fn observe_any(
        &mut self,
        z: &OVector<f64, ZD>,
        h: fn(&OVector<f64, D>) -> OVector<f64, ZD>,
        h_normalize: fn(h: &mut OVector<f64, ZD>, h0: &OVector<f64, ZD>),
        H: &OMatrix<f64, ZD, D>,
        noise: &CorrelatedNoise<f64, ZD>,
    ) -> Result<(), &str> {
        let mut hx = h(&self.state().unwrap());
        h_normalize(&mut hx, &z);
        let s = z - hx;
        ExtendedLinearObserver::observe_innovation(self, &s, H, noise)?;
        Ok(())
    }
}

/// Test information estimator operations defined on a InformationRootState.
impl<D: Dim, QD: Dim, ZD: Dim> FatEstimator<D, QD, ZD> for InformationRootState<f64, D>
where
    // display
    DefaultAllocator: Allocator<D, D>,

    // InformationRootState
    DefaultAllocator: Allocator<D, D> + Allocator<D>,

    // FatEstimator
    DefaultAllocator: Allocator<D, D>
        + Allocator<ZD, D>
        + Allocator<U1, D>
        + Allocator<D>
        + Allocator<QD, QD>
        + Allocator<ZD, ZD>
        + Allocator<D, QD>
        + Allocator<QD>
        + Allocator<ZD>,

    // predict
    D: DimAdd<QD>,
    DefaultAllocator: Allocator<DimSum<D, QD>, DimSum<D, QD>>
        + Allocator<DimSum<D, QD>>
        + Allocator<D, QD>
        + Allocator<QD>,
    DimSum<D, QD>: DimMin<DimSum<D, QD>>,
    DefaultAllocator: Allocator<DimMinimum<DimSum<D, QD>, DimSum<D, QD>>>
        + Allocator<DimMinimum<DimSum<D, QD>, DimSum<D, QD>>, DimSum<D, QD>>,
    // observe
    DefaultAllocator: Allocator<D, D>
        + Allocator<ZD, D>
        + Allocator<ZD, ZD>
        + Allocator<D>
        + Allocator<ZD>,
    D: DimAdd<ZD> + DimAdd<U1>,
    DefaultAllocator: Allocator<DimSum<D, ZD>, DimSum<D, U1>> + Allocator<DimSum<D, ZD>>,
    DimSum<D, ZD>: DimMin<DimSum<D, U1>>,
    DefaultAllocator: Allocator<DimMinimum<DimSum<D, ZD>, DimSum<D, U1>>>
        + Allocator<DimMinimum<DimSum<D, ZD>, DimSum<D, U1>>, DimSum<D, U1>>,
{
    fn dim(&self) -> D {
        return self.r.shape_generic().0;
    }

    fn trace_state(&self) {
        println!("{}", self.R);
    }

    fn init(&mut self, state: &KalmanState<f64, D>) -> Result<(), &str> {
        let info = InformationRootState::try_from(state.clone()).unwrap();
        self.r = info.r;
        self.R = info.R;
        Ok(())
    }

    fn predict_fn(
        &mut self,
        x_pred: &OVector<f64, D>,
        _f: fn(&OVector<f64, D>) -> OVector<f64, D>,
        F: &OMatrix<f64, D, D>,
        noise: &CoupledNoise<f64, D, QD>,
    ) {
        InformationRootState::predict::<QD>(self, x_pred, F, &noise).unwrap();
    }

    fn observe_any(
        &mut self,
        z: &OVector<f64, ZD>,
        h: fn(&OVector<f64, D>) -> OVector<f64, ZD>,
        h_normalize: fn(h: &mut OVector<f64, ZD>, h0: &OVector<f64, ZD>),
        H: &OMatrix<f64, ZD, D>,
        noise: &CorrelatedNoise<f64, ZD>,
    ) -> Result<(), &str>
    where
        DefaultAllocator: Allocator<D, D>
            + Allocator<U1, U1>
            + Allocator<U1, D>
            + Allocator<D, U1>
            + Allocator<D>
            + Allocator<U1>,
    {
        let mut hx = h(&self.state().unwrap());
        h_normalize(&mut hx, &z);
        let s = z - hx;
        ExtendedLinearObserver::observe_innovation(self, &s, H, noise)?;

        Ok(())
    }
}

/// Test UD estimator operations defined on a FatUDState.
pub struct FatUDState<N: RealField, D: Dim>
where
    DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    pub ud: UDState<N, D>,
    pub obs_uncorrelated: bool,
}

impl<N: Copy + RealField, D: Dim> Estimator<N, D> for FatUDState<N, D>
where
    DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    fn state(&self) -> Result<OVector<N, D>, &str> {
        self.ud.state()
    }
}

impl<N: Copy + RealField, D: Dim> KalmanEstimator<N, D> for FatUDState<N, D>
where
    DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    fn kalman_state(&self) -> Result<KalmanState<N, D>, &str> {
        self.ud.kalman_state()
    }
}

impl<D: DimAdd<U1>, QD: Dim, ZD: Dim> FatEstimator<D, QD, ZD> for FatUDState<f64, D>
where
    DefaultAllocator: Allocator<D, D>
        + Allocator<D, D>
        + Allocator<D>
        + Allocator<QD, QD>
        + Allocator<D, QD>
        + Allocator<QD>
        + Allocator<ZD, ZD>
        + Allocator<ZD, D>
        + Allocator<ZD>
        + Allocator<D, DimSum<D, U1>>,
    D: DimAdd<QD>,
    DefaultAllocator: Allocator<DimSum<D, QD>, U1>,
{
    fn trace_state(&self) {
        println!("{}", self.ud.UD);
    }

    fn init(&mut self, state: &KalmanState<f64, D>) -> Result<(), &str> {
        let ud_state = UDState::try_from(state.clone()).unwrap();
        self.ud = ud_state;
        Ok(())
    }

    fn predict_fn(
        &mut self,
        x_pred: &OVector<f64, D>,
        _f: fn(&OVector<f64, D>) -> OVector<f64, D>,
        F: &OMatrix<f64, D, D>,
        noise: &CoupledNoise<f64, D, QD>,
    ) {
        self.ud.predict::<QD>(F, x_pred, noise).unwrap();
    }

    fn observe_any(
        &mut self,
        z: &OVector<f64, ZD>,
        h: fn(&OVector<f64, D>) -> OVector<f64, ZD>,
        h_normalize: fn(h: &mut OVector<f64, ZD>, h0: &OVector<f64, ZD>),
        H: &OMatrix<f64, ZD, D>,
        noise: &CorrelatedNoise<f64, ZD>,
    ) -> Result<(), &str> {
        if self.obs_uncorrelated {
            assert_eq!(noise.Q.upper_triangle() - Matrix::from_diagonal(&noise.Q.diagonal()),
                       Matrix::zeros_generic(noise.Q.shape_generic().0, noise.Q.shape_generic().1));
            let mut hx = h(&self.state().unwrap());
            h_normalize(&mut hx, &z);
            let s = z - &hx;
            let uncorrelated_noise = UncorrelatedNoise::<f64, ZD> {
                q: noise.Q.diagonal(),
            };
            self.ud
                .observe_innovation::<ZD>(&s, H, &uncorrelated_noise)
                .map(|_rcond| {})
        } else {
            let noise_fac = CorrelatedFactorNoise::from_correlated(&noise)?;
            self.ud
                .observe_linear_correlated::<ZD>(z, H, h_normalize, &noise_fac)
                .map(|_rcond| {})
        }
    }
}

/// Test Unscented estimator operations defined on a UnscentedDuplexState.
impl<D: Dim, QD: Dim, ZD: Dim> FatEstimator<D, QD, ZD> for UnscentedDuplexState<f64, D>
where
    DefaultAllocator: Allocator<D, D>
        + Allocator<D>
        + Allocator<QD, QD>
        + Allocator<D, QD>
        + Allocator<QD>
        + Allocator<ZD, ZD>
        + Allocator<ZD, D>
        + Allocator<ZD>,
    // predict_unscented
    DefaultAllocator: Allocator<U1, D>,
    // observe_unscented
    DefaultAllocator: Allocator<D, ZD>,
    DefaultAllocator: Allocator<U1, ZD>,
    //
    DefaultAllocator: Allocator<U1, U1> + Allocator<U1>,
    DefaultAllocator: Allocator<D, U1>,
    DefaultAllocator: Allocator<D, D>,
{
    fn trace_state(&self) {
        println!("{:}", self.kalman.X);
    }

    fn init(&mut self, state: &KalmanState<f64, D>) -> Result<(), &str> {
        let duplex_state = UnscentedDuplexState::try_from(state.clone()).unwrap();
        self.kalman = duplex_state.kalman;
        Ok(())
    }

    fn predict_fn(
        &mut self,
        _x_pred: &OVector<f64, D>,
        f: fn(&OVector<f64, D>) -> OVector<f64, D>,
        _F: &OMatrix<f64, D, D>,
        noise: &CoupledNoise<f64, D, QD>,
    ) {
        self.predict(f, &CorrelatedNoise::from_coupled::<QD>(noise)).unwrap();
    }

    fn observe_any(
        &mut self,
        z: &OVector<f64, ZD>,
        h: fn(&OVector<f64, D>) -> OVector<f64, ZD>,
        h_normalize: fn(h: &mut OVector<f64, ZD>, h0: &OVector<f64, ZD>),
        _H: &OMatrix<f64, ZD, D>,
        noise: &CorrelatedNoise<f64, ZD>,
    ) -> Result<(), &str> {
        let hn = |x: &OVector<f64, D>| -> OVector<f64, ZD> {
            let mut zp = h(x);
            h_normalize(&mut zp, z);
            zp
        };
        self.observe(&z, hn, noise)
    }
}

/// Test SIR estimator operations defined on a FatSampleState.
pub struct FatSampleState<'a, N: RealField, D: Dim>
where
    DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    pub sample: SampleStateDraw<'a, N, D>,
    pub systematic_resampler: bool,
    pub kalman_roughening: bool,
}

impl<'a, N: FromPrimitive + RealField, D: Dim> Estimator<N, D> for FatSampleState<'a, N, D>
where
    DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    fn state(&self) -> Result<OVector<N, D>, &str> {
        self.sample.state.state()
    }
}

impl<'a, N: Copy + FromPrimitive + RealField, D: Dim> KalmanEstimator<N, D> for FatSampleState<'a, N, D>
where
    DefaultAllocator: Allocator<D, D> + Allocator<U1, D> + Allocator<D>,
{
    fn kalman_state(&self) -> Result<KalmanState<N, D>, &str> {
        self.sample.state.kalman_state()
    }
}

impl<'a, D: Dim, QD: Dim, ZD: Dim> FatEstimator<D, QD, ZD> for FatSampleState<'a, f64, D>
where
    DefaultAllocator: Allocator<D, D>
        + Allocator<D>
        + Allocator<QD, QD>
        + Allocator<D, QD>
        + Allocator<QD>
        + Allocator<ZD, ZD>
        + Allocator<ZD, D>
        + Allocator<ZD>,
    // sample
    DefaultAllocator: Allocator<U1, D>,
{
    fn dim(&self) -> D {
        return self.sample.state.s[0].shape_generic().0;
    }

    fn allow_error_by(&self) -> f64 {
        3000000. / (self.sample.state.s.len() as f64).sqrt() // sample error scales with sqrt number samples
    }

    fn trace_state(&self) {
        // println!("{:?}\n{:?}", self.w, self.s);
    }

    fn init(&mut self, state: &KalmanState<f64, D>) -> Result<(), &str> {
        let sample_state = SampleState::from_kalman((*state).clone(), self.sample.state.s.len(), self.sample.rng).unwrap();
        self.sample.state = sample_state;
        Ok(())
    }

    fn predict_fn(
        &mut self,
        _x_pred: &OVector<f64, D>,
        f: fn(&OVector<f64, D>) -> OVector<f64, D>,
        _F: &OMatrix<f64, D, D>,
        noise: &CoupledNoise<f64, D, QD>,
    ) {
        // Predict amd sample the noise
        let cn = &noise.G * Matrix::from_diagonal(&noise.q.map(|v| v.sqrt()));
        let sampler = sir::normal_noise_sampler_coupled(cn);
        self.sample.predict_sampled(
            move |x: &OVector<f64, D>, rng: &mut dyn RngCore| -> OVector<f64, D> {
                f(&x) + sampler(rng)
            },
        );
    }

    fn observe_any(
        &mut self,
        z: &OVector<f64, ZD>,
        h: fn(&OVector<f64, D>) -> OVector<f64, ZD>,
        _h_normalize: fn(h: &mut OVector<f64, ZD>, h0: &OVector<f64, ZD>),
        _H: &OMatrix<f64, ZD, D>,
        noise: &CorrelatedNoise<f64, ZD>,
    ) -> Result<(), &str> {
        self.sample.state.observe(sir::gaussian_observation_likelihood(z, h, noise));

        let mut resampler = if self.systematic_resampler {
            |w: &mut sir::Likelihoods, rng: &mut dyn RngCore| sir::standard_resampler(w, rng)
        } else {
            |w: &mut sir::Likelihoods, rng: &mut dyn RngCore| sir::systematic_resampler(w, rng)
        };

        if self.kalman_roughening {
            let noise = CoupledNoise::from_correlated(&CorrelatedNoise {
                Q: self.kalman_state().unwrap().X,
            })
            .unwrap();
            let mut roughener = move |s: &mut sir::Samples<f64, D>, rng: &mut dyn RngCore| {
                sir::roughen_noise(s, &noise, 1., rng)
            };
            self.sample
                .update_resample(&mut resampler, &mut roughener)?;
        } else {
            let mut roughener = |s: &mut sir::Samples<f64, D>, rng: &mut dyn RngCore| {
                sir::roughen_minmax(s, 1., rng)
            };
            self.sample
                .update_resample(&mut resampler, &mut roughener)?;
        };

        Ok(())
    }
}

trait Diagonal<N: Scalar, R: Dim, C: Dim, S: Storage<N, R, C>> {
    fn diagonal_iter(&self) -> StepBy<MatrixIter<N, R, C, S>>;
}

impl<N: Scalar, R: Dim, C: Dim, S: Storage<N, R, C>> Diagonal<N, R, C, S> for Matrix<N, R, C, S> {
    fn diagonal_iter(&self) -> StepBy<MatrixIter<N, R, C, S>> {
        self.iter().step_by(self.nrows() + 1)
    }
}