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

KalmanFilter::KalmanFilter() {
    // 6 state variables (x, y, vx, vy, ax, ay), 2 measurements (x, y)
    kf.init(6, 2, 0);
    init_params();
    initialized = false;
    has_prev_measurement = false;
    has_prev_velocity = false;
    time_elapsed = 1.0 / 100.0; // Assuming 100 FPS as default
    last_measurement = cv::Point2f(0, 0);
    last_prediction = cv::Point2f(0, 0);
    prev_measurement = cv::Point2f(0, 0);
    prev_velocity = cv::Point2f(0, 0);
}

void KalmanFilter::init_params() {
    // Transition matrix for constant acceleration model
    // x(k+1) = x(k) + vx(k)*dt + 0.5*ax(k)*dt^2
    // y(k+1) = y(k) + vy(k)*dt + 0.5*ay(k)*dt^2
    // vx(k+1) = vx(k) + ax(k)*dt
    // vy(k+1) = vy(k) + ay(k)*dt
    // ax(k+1) = ax(k)
    // ay(k+1) = ay(k)

    double dt = time_elapsed;
    kf.transitionMatrix = (cv::Mat_<float>(6, 6) <<
        1, 0, dt, 0, 0.5*dt*dt, 0,
        0, 1, 0, dt, 0, 0.5*dt*dt,
        0, 0, 1, 0, dt, 0,
        0, 0, 0, 1, 0, dt,
        0, 0, 0, 0, 1, 0,
        0, 0, 0, 0, 0, 1
    );

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

    // Process noise covariance - higher for acceleration components
    kf.processNoiseCov = cv::Mat::zeros(6, 6, CV_32F);
    kf.processNoiseCov.at<float>(0, 0) = 1.0f;      // x position noise
    kf.processNoiseCov.at<float>(1, 1) = 1.0f;      // y position noise
    kf.processNoiseCov.at<float>(2, 2) = 4.0f;      // x velocity noise
    kf.processNoiseCov.at<float>(3, 3) = 4.0f;      // y velocity noise
    kf.processNoiseCov.at<float>(4, 4) = 8.0f;      // x acceleration noise
    kf.processNoiseCov.at<float>(5, 5) = 8.0f;      // y acceleration noise

    // Measurement noise covariance
    kf.measurementNoiseCov = cv::Mat::eye(2, 2, CV_32F) * 0.1f; // Lower measurement noise

    // Error covariance post
    kf.errorCovPost = cv::Mat::eye(6, 6, CV_32F) * 0.1f;
}

cv::Point2f KalmanFilter::estimate_velocity(const cv::Point2f& current, const cv::Point2f& previous) {
    // Calculate velocity based on position difference
    cv::Point2f velocity = current - previous;
    velocity.x /= (float)time_elapsed;
    velocity.y /= (float)time_elapsed;
    return velocity;
}

void KalmanFilter::updateProcessNoise(const cv::Point2f& current_measurement) {
    if (has_prev_measurement && has_prev_velocity) {
        // Estimate current velocity
        cv::Point2f current_velocity = estimate_velocity(current_measurement, prev_measurement);

        // Calculate acceleration based on change in velocity
        cv::Point2f acceleration = current_velocity - prev_velocity;
        acceleration.x /= (float)time_elapsed;
        acceleration.y /= (float)time_elapsed;

        // Adjust process noise based on estimated acceleration (higher acceleration = higher uncertainty)
        float acc_magnitude = std::sqrt(acceleration.x * acceleration.x + acceleration.y * acceleration.y);

        // Increase process noise for rapid maneuvers
        float noise_factor = 1.0f + 0.5f * acc_magnitude;

        kf.processNoiseCov.at<float>(4, 4) = 8.0f * noise_factor;  // ax noise
        kf.processNoiseCov.at<float>(5, 5) = 8.0f * noise_factor;  // ay noise
    }
}

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>(6, 1) <<
            measurement.x,         // x position
            measurement.y,         // y position
            0.0f,                  // x velocity
            0.0f,                  // y velocity
            0.0f,                  // x acceleration
            0.0f                   // y acceleration
        );
        initialized = true;
    } else {
        // Update process noise based on target movement characteristics
        updateProcessNoise(measurement);

        // Perform correction step
        kf.correct(measurement_mat);
    }

    // Update previous state for next iteration
    if (has_prev_measurement) {
        // Calculate velocity based on position change
        cv::Point2f current_velocity = estimate_velocity(measurement, prev_measurement);

        if (has_prev_velocity) {
            // Estimate acceleration based on velocity change
            cv::Point2f acceleration = current_velocity - prev_velocity;
            acceleration.x /= (float)time_elapsed;
            acceleration.y /= (float)time_elapsed;

            // Update state with estimated acceleration
            kf.statePost.at<float>(4) = acceleration.x; // x acceleration
            kf.statePost.at<float>(5) = acceleration.y; // y acceleration
        }

        // Update previous velocity
        prev_velocity = current_velocity;
        has_prev_velocity = true;
    } else {
        // Initialize previous velocity as zero
        prev_velocity = cv::Point2f(0, 0);
        has_prev_velocity = true;
    }

    // Update stored measurements
    prev_measurement = last_measurement;
    has_prev_measurement = true;
    last_measurement = measurement;
}

cv::Point2f KalmanFilter::get_current_velocity() const {
    if (!initialized) {
        return cv::Point2f(0.0f, 0.0f);
    }

    // Return the velocity components from the state vector
    float vx = kf.statePost.at<float>(2);
    float vy = kf.statePost.at<float>(3);
    return cv::Point2f(vx, vy);
}

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;
}