/*
 * KalmanWrapper.cpp
 *
 * A wrapper for the OpenCV KalmanFilter class which offers a couple of
 * extra functionalities and default parameters.
 *
 *  Created on: Apr 3, 2020
 *      Author: lukas
 */

#include "KalmanWrapper.hpp"
#include "opencv2/opencv.hpp"

using namespace cv;
using namespace std;

/**
 * Wrapper for a KalmanFilter class
 *
 * Creates a Kalmfilter class with the default parameters from the lecture.
 *
 * @param model_type String, the type of the model. Either "constant_velocity" or
 * "constant_acceleration".
 */
KalmanWrapper::KalmanWrapper(String model_type) {
	_model_type = model_type;
	_meas_size = 2;

	_default_R = Mat::eye(_meas_size, _meas_size, CV_32F);
	_default_R *= 25;

	if (_model_type == "constant_velocity")
	{
		_state_size = 4;

		_A = Mat::eye(_state_size, _state_size, CV_32F);
		_A.at<float>(0,1) = 1;
		_A.at<float>(2,3) = 1;

		_H = Mat::zeros(_meas_size, _state_size, CV_32F);
		_H.at<float>(0,0) = 1;
		_H.at<float>(1,2) = 1;

		_default_Q = Mat::zeros(_state_size, _state_size, CV_32F);
		_default_Q.at<float>(0,0) = 25;
		_default_Q.at<float>(1,1) = 10;
		_default_Q.at<float>(2,2) = 25;
		_default_Q.at<float>(3,3) = 10;

		_default_P = Mat::eye(_state_size, _state_size, CV_32F);
		_default_P *= 100000;
	}
	else if (_model_type == "constant_acceleration")
	{
		_state_size = 6;

		_A = Mat::eye(_state_size, _state_size, CV_32F);
		_A.at<float>(0,1) = 1;
		_A.at<float>(0,2) = 0.5;
		_A.at<float>(1,2) = 1;
		_A.at<float>(3,4) = 1;
		_A.at<float>(3,5) = 0.5;
		_A.at<float>(4,5) = 1;

		_H = Mat::zeros(_meas_size, _state_size, CV_32F);
		_H.at<float>(0,0) = 1;
		_H.at<float>(1,3) = 1;

		_default_Q = Mat::zeros(_state_size, _state_size, CV_32F);
		_default_Q.at<float>(0,0) = 25;
		_default_Q.at<float>(1,1) = 10;
		_default_Q.at<float>(2,2) = 1;
		_default_Q.at<float>(3,3) = 25;
		_default_Q.at<float>(4,4) = 10;
		_default_Q.at<float>(5,5) = 1;

		_default_P = Mat::eye(_state_size, _state_size, CV_32F);
		_default_P *= 100000;
	}
	else
	{
		cout << "Model type not supported" << endl;
	}

	_default_P.copyTo(_P);
	_default_Q.copyTo(_Q);
	_default_R.copyTo(_R);

	_kf = KalmanFilter(_state_size, _meas_size, 0, CV_32F);
	_kf.transitionMatrix = _A;
	_kf.measurementMatrix = _H;
	_kf.errorCovPost = _P;
	_kf.processNoiseCov = _Q;
	_kf.measurementNoiseCov = _R;
}

KalmanWrapper::~KalmanWrapper() {
	// TODO Auto-generated destructor stub
}

/**
 * Generates a prediction of with the Kalman filter.
 *
 * When generating a prediction, the passed pointer variable is
 * populated with the x and y value of the prediction, as the velocity/acceleration
 * is usually not of interest for the client.
 *
 * @param &prediction Mat, row vector of size 2 which will be populated
 */
void KalmanWrapper::generatePrediction(Mat &prediction) {
	Mat predicted_Result = _kf.predict();

	if (_model_type == "constant_velocity")
	{
		prediction.at<float>(0,0) = predicted_Result.at<float>(0,0);
		prediction.at<float>(1,0) = predicted_Result.at<float>(2,0);
	}
	else if (_model_type == "constant_acceleration")
	{
		prediction.at<float>(0,0) = predicted_Result.at<float>(0,0);
		prediction.at<float>(1,0) = predicted_Result.at<float>(3,0);
	}
}

/**
 * Generates a corrected measurement with the Kalman filter.
 *
 * When generating a corrected measurement, the passed pointer variable is
 * populated with the x and y value of the corrected measurement, as the velocity/acceleration
 * is usually not of interest for the client.
 *
 * @param measurement Mat, row vector of size 2 containing the current observation.
 * @param &corrected_measurement Mat, row vector of size 2 which will be populated.
 */
void KalmanWrapper::generateCorrection(Mat measurement, Mat &corrected_measurement) {
	Mat corrected_Result = _kf.correct(measurement);
//	cout << "Gain:" << endl;
//	cout << _kf.gain << endl << endl;

	if (_model_type == "constant_velocity")
	{
		corrected_measurement.at<float>(0,0) = corrected_Result.at<float>(0,0);
		corrected_measurement.at<float>(1,0) = corrected_Result.at<float>(2,0);
	}
	else if (_model_type == "constant_acceleration")
	{
		corrected_measurement.at<float>(0,0) = corrected_Result.at<float>(0,0);
		corrected_measurement.at<float>(1,0) = corrected_Result.at<float>(3,0);
	}
}

/**
 * Sets the x and y state of the Kalman filter.
 *
 * Helper function that simplifies setting the x and y state of the Kalman filter,
 * independently from the model type.
 *
 * @param observation Mat, row vector of size 2 containing the observation.
 */
void KalmanWrapper::setInitialXY(Mat observation) {
	int y_pos = _state_size / 2;

	Mat initial_state = Mat::zeros(_state_size,1, CV_32F);
	initial_state.at<float>(0,0) = observation.at<float>(0,0);
	initial_state.at<float>(y_pos,0) = observation.at<float>(1,0);

	setInitialState(initial_state);
}

/**
 * Sets the parameters for a specific sequence.
 *
 * Given a path or video name, the parameters that are empirically best for the
 * sequence in question are set
 *
 * @param path String, the path or name of the video
 */
void KalmanWrapper::setParams(string path) {
	string vid_name = path.substr(path.rfind("/") + 1);

	// Set parameters to default and then mask changes corresponding
	// to the video name
	_default_R.copyTo(_R);
	_default_P.copyTo(_P);
	_default_Q.copyTo(_Q);

	cout << vid_name << endl;

	if (vid_name == "video6.mp4")
	{
		if(_model_type == "constant_acceleration")
		{
			_R.at<float>(1,1) = 2500;

			_Q.at<float>(2,2) = 0.001;
			_Q.at<float>(4,4) = 0.001;
			_Q.at<float>(5,5) = 0.001;

			_P.at<float>(4,4) = 10;
			_P.at<float>(5,5) = 10;

		} else {
			_Q.at<float>(3,3) = 0.001;
		}
	}

	else if (vid_name == "abandonedBox_600_1000_clip.mp4")
	{
		// We trust the observation most
		_R.at<float>(0,0) = 2.5;
		_R.at<float>(1,1) = 2.5;
	}

	else if (vid_name == "boats_6950_7900_clip.mp4")
	{
		if(_model_type == "constant_acceleration")
		{
			// Measurement is noisy along the y-axis
			_R.at<float>(1,1) = 2500;

			// No process noise along the y-axis
			_Q.at<float>(4,4) = 0.001;
			_Q.at<float>(5,5) = 0.0001;

			// The transition model is relatively stabel along the y-axis
			_P.at<float>(4,4) = 10;
			_P.at<float>(5,5) = 10;

		}
		else
		{
			// Measurement is noisy along the y-axis
			_R.at<float>(1,1) = 2500;

			// No process noise along the y-axis
			_Q.at<float>(3,3) = 0.001;

			// The transition model is relatively stabel along the y-axis
			_P.at<float>(3,3) = 10;

		}

	}

	else if (vid_name == "pedestrians_800_1025_clip.mp4")
	{
		if(_model_type == "constant_acceleration"){
			// Measurement is noisy along the y-axis
			_R.at<float>(1,1) = 25000;
			_R.at<float>(0,0) = 25000;

			// No process noise along the y-axis
			_Q.at<float>(1,1) = 0.001;
			_Q.at<float>(2,2) = 0.0001;
			_Q.at<float>(4,4) = 0.001;
			_Q.at<float>(5,5) = 0.0001;
		}
		else
		{
			// Measurement is noisy along the y-axis
			_R.at<float>(1,1) = 25000;
			_R.at<float>(0,0) = 25000;

			// Process is stable when it comes to velocity
			_Q.at<float>(1,1) = 0.001;
			_Q.at<float>(3,3) = 0.001;

		}
	}

	// If no params found for video
	// Set default params for video
	else
	{
		cout<<"Video parameters not found!"<<endl;
		cout<<"using default parameters!"<<endl;
	}

	_kf.errorCovPost = _P;
	_kf.processNoiseCov = _Q;
	_kf.measurementNoiseCov = _R;
}