#include "BallisticPredictor.h"
#include "config.h"
#include <cmath>

BallisticPredictor::BallisticPredictor(float bullet_speed) 
    : bullet_speed(bullet_speed), last_ballistic_point(nullptr) {
    armor_half_width = ARMOR_WIDTH / 2000.0f;  // Convert to meters
    armor_half_height = ARMOR_HEIGHT / 2000.0f;  // Convert to meters
}

cv::Point2f* BallisticPredictor::predict_ballistic_point(const cv::Point2f* predicted_center, 
                                                        const cv::Point2f& img_center, 
                                                        const cv::Point2f& target_speed) {
    if (predicted_center == nullptr) {
        if (last_ballistic_point != nullptr) {
            delete last_ballistic_point;
            last_ballistic_point = nullptr;
        }
        return nullptr;
    }

    // 1. Estimate target distance (meters)
    float distance = calculate_distance(*predicted_center, img_center);

    // 2. Calculate bullet flight time (seconds): distance / bullet speed
    float flight_time = bullet_speed > 0 ? distance / bullet_speed : 0.0f;

    // 3. Calculate target movement in flight time (pixels)
    // target speed (pixels/s) * flight time (s) = predicted movement in pixels
    float delta_x = target_speed.x * flight_time;
    float delta_y = target_speed.y * flight_time;

    // 4. Ballistic prediction point = Kalman prediction center + target movement compensation
    float ballistic_x = predicted_center->x + delta_x;
    float ballistic_y = predicted_center->y + delta_y;
    cv::Point2f ballistic_point(ballistic_x, ballistic_y);

    // Update last ballistic point
    if (last_ballistic_point != nullptr) {
        delete last_ballistic_point;
    }
    last_ballistic_point = new cv::Point2f(ballistic_point);

    return last_ballistic_point;
}

float BallisticPredictor::calculate_distance(const cv::Point2f& armor_center, const cv::Point2f& img_center, float focal_length) {
    // Calculate the pixel distance from armor center to image center
    float dx = std::abs(armor_center.x - img_center.x);
    float dy = std::abs(armor_center.y - img_center.y);
    float pixel_distance = std::sqrt(dx * dx + dy * dy);

    // If target is near the center, avoid division by zero
    if (pixel_distance < 10) {
        return 5.0f;  // Default 5 meters (adjust based on actual scenario)
    }

    // Distance estimation formula: (actual size * focal length) / pixel size
    float distance = (armor_half_width * focal_length) / (pixel_distance * 0.5f);
    return std::max(0.5f, distance);  // Limit minimum distance to 0.5 meters
}