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 measurementNoiseFactor, processNoiseFactor;

        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
        public Kalman2DPositionFilter(float measurementNoiseFactor = 1.0e-1f, float processNoiseFactor = 1.0e-4f)
        {
            this.measurementNoiseFactor = measurementNoiseFactor;
            this.processNoiseFactor = processNoiseFactor;
            reset();
        }

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

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

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

            // measurement noise covariance matrix
            Matrix<float> measurementNoiseCovarianceMatrix = new Matrix<float>(2, 2);
            measurementNoiseCovarianceMatrix.SetIdentity(new MCvScalar(measurementNoiseFactor));
            kalman.MeasurementNoiseCovariance = measurementNoiseCovarianceMatrix;

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

            // error covariance post matrix (initial value)
            Matrix<float> errorCovariancePostMatrix = new Matrix<float>(4, 4);
            errorCovariancePostMatrix.SetIdentity();
            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, 0.0f, 0.0f });
            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
            Matrix<float> prediction = kalman.Predict();

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