using System;
using System.Collections.Generic;
using System.Text;
using Mapack;

namespace Kalman
{
    class KalmanFilter
    {
        Matrix X = null;    // State Variables
        Matrix Z = null;    // Output/Sensors
        Matrix U = null;    // Control Vector

        Matrix B = null;    // Control Matrix

        Matrix F = null;    // State Update Matrix
        Matrix H = null;    // Output Matrix

        Matrix P = null;    // Covariance Matrix

        Matrix R = null;    // Measurement Noise Covariance
        Matrix Q = null;    // System Noise Covariance

        public KalmanFilter(int StateVariableCount, int OutputCount, int ControlCount)
        {
            X = new Matrix(StateVariableCount, 1);
            Z = new Matrix(OutputCount, 1);

            U = new Matrix(ControlCount, 1);
            B = new Matrix(StateVariableCount, ControlCount);

            P = new Matrix(StateVariableCount, StateVariableCount);

            R = new Matrix(OutputCount, OutputCount);
            Q = new Matrix(StateVariableCount, StateVariableCount);
        }

        public void SetOutput(int i, double value)
        {
            Z[i, 0] = value;
        }

        public void SetControl(int i, double value)
        {
            U[i, 0] = value;
        }

        public void Predict()
        {
            X = F * X + B * U;  //Predicted state
            P = F * P * F.Transpose() + Q;  //Predicted estimate covariance
        }

        public void Update()
        {
            Matrix Y = Z - H * X;   //Measurement residual
            Matrix S = H * P * H.Transpose() + R;   //Residual Covariance
            Matrix K = P * H.Transpose() * S.Inverse;   //Kalman Gain
            X = X + K * Y;  //Updated state estimate;
            P = (Matrix.Diagonal(P.Rows, P.Columns, 1) - K * H) * P;    //updated estimate covariance
        }

        public Matrix StateUpdateMatrix
        {
            set
            {
                F = value;
            }
        }

        public void SetProcessNoise(double[] noise)
        {
            Matrix Noise = new Matrix(new double[][] { noise });
            Q = Noise.Transpose() * Noise;
        }

        public void SetMeasurementNoise(double[] noise)
        {
            Matrix Noise = new Matrix(new double[][] { noise });
            R = Noise.Transpose() * Noise;
        }

        public void SetInitialConfidence(double Confidence)
        {
            //Initialize the diagonal of the Covariance Matrix with a
            //value representing confidence in the initial position
            Q = Matrix.Diagonal(Q.Rows, Q.Columns, Confidence);
        }

        public double this[int i]
        {
            get
            {
                return X[i, 0];
            }

            set
            {
                X[i, 0] = value;
            }
        }

        public Matrix OutputMatrix
        {
            set
            {
                H = value;
            }
        }

        public Matrix ControlMatrix
        {
            set
            {
                B = value;
            }
        }
    }
}