#![allow(non_snake_case)]

//! Information 'square root' state estimation.
//!
//! A discrete Bayesian estimator that uses a linear information root representation [`InformationRootState`] of the system for estimation.
//!
//! The linear representation can also be used for non-linear systems by using linearised models.

use nalgebra::{
    allocator::Allocator, Const, DefaultAllocator, Dim, DimAdd, DimMin, DimMinimum, DimSum, U1,
};
use nalgebra::{OMatrix, OVector, RealField, QR};

use crate::linalg::cholesky::UDU;
use crate::matrix::check_positive;
use crate::models::{
    Estimator, ExtendedLinearObserver, InformationState, KalmanEstimator, KalmanState,
};
use crate::noise::{CorrelatedNoise, CoupledNoise};

/// Information State.
///
/// Linear representation as a information root state vector and the information root (upper triangular) matrix.
/// For a given [KalmanState] the information root state inverse(R).inverse(R)' == X, r == R.x
/// For a given [InformationState] the information root state R'.R == I, r == inverse(R).i
#[derive(PartialEq, Clone)]
pub struct InformationRootState<N: RealField, D: Dim>
where
    DefaultAllocator: Allocator<N, D, D> + Allocator<N, D>,
{
    /// Information root state vector
    pub r: OVector<N, D>,
    /// Information root matrix (upper triangular)
    pub R: OMatrix<N, D, D>,
}

impl<N: Copy + RealField, D: Dim> TryFrom<InformationState<N, D>> for InformationRootState<N, D>
where
    DefaultAllocator: Allocator<N, D, D> + Allocator<N, D>,
{
    type Error = &'static str;

    fn try_from(state: InformationState<N, D>) -> Result<Self, Self::Error> {
        // Information Root, R'.R = I
        let R = state.I.clone().cholesky().ok_or("I not PD")?
            .l().transpose();

        // Information Root state, r=inv(R)'.i
        let shape = R.shape_generic();
        let mut RI = OMatrix::identity_generic(shape.0, shape.1);
        R.solve_upper_triangular_mut(&mut RI);
        let r = RI.tr_mul(&state.i);
        Ok(InformationRootState { r, R })
    }
}

impl<N: Copy + RealField, D: Dim> InformationRootState<N, D>
    where
        DefaultAllocator: Allocator<N, D, D> + Allocator<N, D>,
{
    pub fn information_state(&self) -> Result<InformationState<N, D>, &str> {
        // Information, I = R.R'
        let I = self.R.tr_mul(&self.R);
        let x = self.state()?;
        // Information state, i = I.x
        let i = &I * x;

        Ok(InformationState { i, I })
    }
}

impl<N: Copy + RealField, D: Dim> Estimator<N, D> for InformationRootState<N, D>
where
    DefaultAllocator: Allocator<N, D, D> + Allocator<N, D>,
{
    fn state(&self) -> Result<OVector<N, D>, &str> {
        self.kalman_state().map(|res| res.x)
    }
}

impl<N: Copy + RealField, D: Dim> TryFrom<KalmanState<N, D>> for InformationRootState<N, D>
    where
        DefaultAllocator: Allocator<N, D, D> + Allocator<N, D>,
{
    type Error = &'static str;

    fn try_from(state: KalmanState<N, D>) -> Result<Self, Self::Error> {
        // Information Root, inv(R).inv(R)' = X
        let udu = UDU::new();
        let mut R = state.X.clone();
        let rcond = udu.UCfactor_n(&mut R, state.X.nrows());
        check_positive(rcond, "X not PD")?;
        udu.UTinverse(&mut R).unwrap(); // check_positive should prevent singular

        // Information Root state, r=R*x
        let r = &R * &state.x;
        Ok( InformationRootState { r, R } )
    }
}

impl<N: Copy + RealField, D: Dim> KalmanEstimator<N, D> for InformationRootState<N, D>
where
    DefaultAllocator: Allocator<N, D, D> + Allocator<N, D>,
{
    fn kalman_state(&self) -> Result<KalmanState<N, D>, &str> {
        let shape = self.R.shape_generic();
        let mut RI = OMatrix::identity_generic(shape.0, shape.1);
        self.R.solve_upper_triangular_mut(&mut RI);

        // Covariance X = inv(R).inv(R)'
        let X = &RI * &RI.transpose();
        // State, x= inv(R).r
        let x = RI * &self.r;

        Ok(KalmanState { x, X })
    }
}

impl<N: Copy + RealField, D: Dim, ZD: Dim> ExtendedLinearObserver<N, D, ZD>
    for InformationRootState<N, D>
where
    DefaultAllocator: Allocator<N, D, D>
        + Allocator<N, ZD, D>
        + Allocator<N, ZD, ZD>
        + Allocator<N, D>
        + Allocator<N, ZD>,
    D: DimAdd<ZD> + DimAdd<U1>,
    DefaultAllocator: Allocator<N, DimSum<D, ZD>, DimSum<D, U1>> + Allocator<N, DimSum<D, ZD>>,
    DimSum<D, ZD>: DimMin<DimSum<D, U1>>,
    DefaultAllocator: Allocator<N, DimMinimum<DimSum<D, ZD>, DimSum<D, U1>>>
        + Allocator<N, DimMinimum<DimSum<D, ZD>, DimSum<D, U1>>, DimSum<D, U1>>,
{
    fn observe_innovation(
        &mut self,
        s: &OVector<N, ZD>,
        hx: &OMatrix<N, ZD, D>,
        noise: &CorrelatedNoise<N, ZD>,
    ) -> Result<(), &str> {
        let udu = UDU::new();
        let mut QI = noise.Q.clone();
        let rcond = udu.UCfactor_n(&mut QI, s.nrows());
        check_positive(rcond, "Q not PD")?;
        udu.UTinverse(&mut QI).unwrap(); // check_positive should prevent singular

        let x = self.state().unwrap(); // state is always defined
        self.observe_info(&(s + hx * x), hx, &QI)
    }
}

impl<N: Copy + RealField, D: Dim> InformationRootState<N, D>
where
    DefaultAllocator: Allocator<N, D, D> + Allocator<N, D>,
{
    pub fn predict<QD: Dim>(
        &mut self,
        x_pred: &OVector<N, D>,
        fx: &OMatrix<N, D, D>,
        noise: &CoupledNoise<N, D, QD>,
    ) -> Result<(), &str>
    where
        D: DimAdd<QD>,
        DefaultAllocator: Allocator<N, DimSum<D, QD>, DimSum<D, QD>>
            + Allocator<N, DimSum<D, QD>>
            + Allocator<N, D, QD>
            + Allocator<N, QD>,
        DimSum<D, QD>: DimMin<DimSum<D, QD>>,
        DefaultAllocator: Allocator<N, DimMinimum<DimSum<D, QD>, DimSum<D, QD>>>
            + Allocator<N, DimMinimum<DimSum<D, QD>, DimSum<D, QD>>, DimSum<D, QD>>,
    {
        let mut Fx_inv = fx.clone();
        let invertable = Fx_inv.try_inverse_mut();
        if !invertable {
            return Err("Fx not invertable")?;
        }

        self.predict_inv_model(x_pred, &Fx_inv, noise)
            .map(|_rcond| {})
    }

    pub fn predict_inv_model<QD: Dim>(
        &mut self,
        x_pred: &OVector<N, D>,
        fx_inv: &OMatrix<N, D, D>, // Inverse of linear prediction model Fx
        noise: &CoupledNoise<N, D, QD>,
    ) -> Result<N, &str>
    where
        D: DimAdd<QD>,
        DefaultAllocator: Allocator<N, DimSum<D, QD>, DimSum<D, QD>>
            + Allocator<N, DimSum<D, QD>>
            + Allocator<N, D, QD>
            + Allocator<N, QD>,
        DimSum<D, QD>: DimMin<DimSum<D, QD>>,
        DefaultAllocator: Allocator<N, DimMinimum<DimSum<D, QD>, DimSum<D, QD>>>
            + Allocator<N, DimMinimum<DimSum<D, QD>, DimSum<D, QD>>, DimSum<D, QD>>,
    {
        // Require Root of correlated predict noise (may be semi-definite)
        let mut Gqr = noise.G.clone();

        for qi in 0..noise.q.nrows() {
            let mut ZZ = Gqr.column_mut(qi);
            ZZ *= noise.q[qi].sqrt();
        }

        // Form Augmented matrix for factorisation
        let dqd = noise.G.shape_generic().0.add(noise.q.shape_generic().0);
        let mut A = OMatrix::identity_generic(dqd, dqd); // Prefill with identity for top left and zero's in off diagonals
        let RFxI: OMatrix<N, D, D> = &self.R * fx_inv;
        let x: OMatrix<N, D, QD> = &RFxI * &Gqr;
        let x_size = x_pred.shape_generic().0;
        let q_size = noise.q.shape_generic().0;
        A.generic_view_mut((q_size.value(), 0), (x_size, q_size))
            .copy_from(&x);
        A.generic_view_mut((q_size.value(), q_size.value()), (x_size, x_size))
            .copy_from(&RFxI);
        A.generic_view_mut((q_size.value(), 0), (x_size, q_size))
            .copy_from(&x);
        A.generic_view_mut((q_size.value(), q_size.value()), (x_size, x_size))
            .copy_from(&RFxI);

        // Calculate factorisation so we have and upper triangular R
        let qr = QR::new(A);
        // Extract the roots
        let r = qr.r();
        self.R
            .copy_from(&r.generic_view((q_size.value(), q_size.value()), (x_size, x_size)));

        self.r = &self.R * x_pred; // compute r from x_pred

        return Ok(UDU::new().UCrcond(&self.R)); // compute rcond of result
    }

    pub fn observe_info<ZD: Dim>(
        &mut self,
        z: &OVector<N, ZD>,
        hx: &OMatrix<N, ZD, D>,
        noise_inv: &OMatrix<N, ZD, ZD>, // Inverse of correlated noise model
    ) -> Result<(), &str>
    where
        DefaultAllocator: Allocator<N, D, D>
            + Allocator<N, ZD, D>
            + Allocator<N, ZD, ZD>
            + Allocator<N, D>
            + Allocator<N, ZD>,
        D: DimAdd<ZD> + DimAdd<U1>,
        DefaultAllocator: Allocator<N, DimSum<D, ZD>, DimSum<D, U1>> + Allocator<N, DimSum<D, ZD>>,
        DimSum<D, ZD>: DimMin<DimSum<D, U1>>,
        DefaultAllocator: Allocator<N, DimMinimum<DimSum<D, ZD>, DimSum<D, U1>>>
            + Allocator<N, DimMinimum<DimSum<D, ZD>, DimSum<D, U1>>, DimSum<D, U1>>,
    {
        let x_size = self.r.shape_generic().0;
        let z_size = z.shape_generic().0;
        // Size consistency, z to model
        if z_size != hx.shape_generic().0 {
            return Err("observation and model size inconsistent");
        }

        // Form augmented matrix for factorisation
        let xd = self.r.shape_generic().0;
        // Prefill with identity for top left and zero's in off diagonals
        let mut A = OMatrix::identity_generic(xd.add(z.shape_generic().0), xd.add(Const::<1>));
        A.generic_view_mut((0, 0), (x_size, x_size))
            .copy_from(&self.R);
        A.generic_view_mut((0, x_size.value()), (x_size, Const::<1>))
            .copy_from(&self.r);
        A.generic_view_mut((x_size.value(), 0), (z_size, x_size))
            .copy_from(&(noise_inv * hx));
        A.generic_view_mut((x_size.value(), x_size.value()), (z_size, Const::<1>))
            .copy_from(&(noise_inv * z));

        // Calculate factorisation so we have and upper triangular R
        let qr = QR::new(A);
        // Extract the roots
        let r = qr.r();

        // Extract the roots
        self.R.copy_from(&r.generic_view((0, 0), (x_size, x_size)));
        self.r
            .copy_from(&r.generic_view((0, x_size.value()), (x_size, Const::<1>)));

        Ok(())
    }
}