#include "KalmanFilter.h"
#include "config.h"
#include <iostream>

KalmanFilter::KalmanFilter() {
    // 4 state variables (x, y, dx, dy), 2 measurements (x, y)
    kf.init(4, 2, 0);
    init_params();
    initialized = false;
    last_measurement = cv::Point2f(0, 0);
    last_prediction = cv::Point2f(0, 0);
}

void KalmanFilter::init_params() {
    // Transition matrix: x(k+1) = x(k) + dx(k), y(k+1) = y(k) + dy(k), dx(k+1) = dx(k), dy(k+1) = dy(k)
    kf.transitionMatrix = (cv::Mat_<float>(4, 4) <<
        1, 0, 1, 0,
        0, 1, 0, 1,
        0, 0, 1, 0,
        0, 0, 0, 1
    );

    // Measurement matrix: measuring x and y
    kf.measurementMatrix = (cv::Mat_<float>(2, 4) <<
        1, 0, 0, 0,
        0, 1, 0, 0
    );

    // Process noise covariance
    kf.processNoiseCov = cv::Mat::eye(4, 4, CV_32F) * KF_PROCESS_NOISE;

    // Measurement noise covariance
    kf.measurementNoiseCov = cv::Mat::eye(2, 2, CV_32F) * KF_MEASUREMENT_NOISE;

    // Error covariance post
    kf.errorCovPost = cv::Mat::eye(4, 4, CV_32F);
}

void KalmanFilter::update(const cv::Point2f& measurement) {
    cv::Mat measurement_mat = (cv::Mat_<float>(2, 1) << measurement.x, measurement.y);

    if (!initialized) {
        // Initialize state if not done yet
        kf.statePost = (cv::Mat_<float>(4, 1) << 
            measurement.x, 
            measurement.y, 
            0.0f, 
            0.0f);
        initialized = true;
    } else {
        kf.correct(measurement_mat);
    }

    last_measurement = measurement;
}

cv::Point2f KalmanFilter::predict() {
    if (!initialized) {
        return cv::Point2f(0, 0);
    }

    cv::Mat prediction = kf.predict();
    cv::Point2f result(prediction.at<float>(0), prediction.at<float>(1));
    last_prediction = result;
    return result;
}