using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

using Emgu.CV;
using Emgu.CV.Structure;

namespace bbiwarg.Utility
{
    class Kalman2DPositionFilter
    {
        private Kalman kalman;
        private float mXX, mXY, mYY, processNoiseFactor;
        private float fps;

        public bool Initialized { get; private set; }

        // state: x, y, v_x, v_y
        //   x: current x position
        //   y: current y position
        // v_x: velocity in x direction
        // v_y: velocity in y direction
        // a_x: acceleration in x direction
        // a_y: acceleration in y direction
        public Kalman2DPositionFilter(float mXX, float mXY, float mYY, float processNoiseFactor = 1.0e-4f, int fps = 30)
        {
            this.mXX = mXX;
            this.mXY = mXY;
            this.mYY = mYY;
            this.processNoiseFactor = processNoiseFactor;
            this.fps = fps;
            reset();
        }

        public void reset()
        {
            // 6 state variables and 2 measurements (0 controls)
            kalman = new Kalman(2, 2, 0);

            // time step (s)
            float t = 1 / fps;

            // transition matrix 
            Matrix<float> transitionMatrix = new Matrix<float>(new float[,] 
            { {1.0f, 0.0f},  
              {0.0f, 1.0f}});
            kalman.TransitionMatrix = transitionMatrix;

            // measurement matrix
            Matrix<float> measurementMatrix = new Matrix<float>(new float[,]
            { {1.0f, 0.0f}, // first measurement = x
              {0.0f, 1.0f}  // second measurement = y
            });
            kalman.MeasurementMatrix = measurementMatrix;

            // measurement noise covariance matrix
            Matrix<float> measurementNoiseCovarianceMatrix = new Matrix<float>(2, 2);

            measurementNoiseCovarianceMatrix[0, 0] = mXX;
            measurementNoiseCovarianceMatrix[0, 1] = measurementNoiseCovarianceMatrix[1, 0] = mXY;
            measurementNoiseCovarianceMatrix[1, 1] = mYY;

            kalman.MeasurementNoiseCovariance = measurementNoiseCovarianceMatrix;

            // process noise covariance matrix
            Matrix<float> processNoiseCovarianceMatrix = new Matrix<float>(2, 2);
            processNoiseCovarianceMatrix.SetIdentity(new MCvScalar(processNoiseFactor));
            kalman.ProcessNoiseCovariance = processNoiseCovarianceMatrix;

            // error covariance post matrix (initial value)
            Matrix<float> errorCovariancePostMatrix = new Matrix<float>(2, 2);
            errorCovariancePostMatrix.SetIdentity(new MCvScalar(processNoiseFactor));
            kalman.ErrorCovariancePost = errorCovariancePostMatrix;

            Initialized = false;
        }

        public void setInitialPosition(Vector2D initialPosition)
        {
            // initial state (x, y, v_x, v_y)
            Matrix<float> initialState = new Matrix<float>(new float[] { initialPosition.X, initialPosition.Y});
            kalman.CorrectedState = initialState;

            Initialized = true;
        }

        public Vector2D getPrediction()
        {
            Matrix<float> predicton = kalman.Predict();
            return new Vector2D(predicton[0, 0], predicton[1, 0]);
        }

        // updates the filter and returns the corrected position
        public Vector2D getCorrectedPosition(Vector2D rawPosition)
        {
            Matrix<float> rawPositionMatrix = new Matrix<float>(new float[,] 
            { {rawPosition.X},
              {rawPosition.Y}});

            // prediction according to model
            kalman.Predict();

            // corrected point
            Matrix<float> estimate = kalman.Correct(rawPositionMatrix);
            return new Vector2D(estimate[0, 0], estimate[1, 0]);
        }
    }
}