import numpy as np
from config import ARMOR_WIDTH, ARMOR_HEIGHT  # 装甲板尺寸（毫米）


class BallisticPredictor:
    """弹道预测器：结合卡尔曼预测位置、距离和弹速计算击中预测点"""

    def __init__(self, bullet_speed=16.0):  # 出弹速度默认16m/s，可调节
        self.bullet_speed = bullet_speed  # 子弹速度（米/秒）
        self.armor_half_width = ARMOR_WIDTH / 2000.0  # 装甲板半宽（转换为米）
        self.armor_half_height = ARMOR_HEIGHT / 2000.0  # 装甲板半高（转换为米）
        self.last_ballistic_point = None  # 记录上一次弹道预测点

    def _calculate_distance(self, armor_center, img_center, focal_length=800):
        """
        估算目标距离（简化模型，需根据实际相机内参调整）
        原理：基于装甲板实际尺寸和图像中尺寸的比例关系
        """
        # 假设图像中装甲板宽度约为灯条间距（简化计算，实际可根据检测的灯条尺寸优化）
        # 此处使用中心偏移量和焦距估算，更精确的方式需结合实际检测的装甲板像素尺寸
        dx = abs(armor_center[0] - img_center[0])
        dy = abs(armor_center[1] - img_center[1])
        pixel_distance = np.sqrt(dx ** 2 + dy ** 2)

        # 若目标在中心附近，避免除以零
        if pixel_distance < 10:
            return 5.0  # 默认5米（可根据实际场景调整）

        # 距离估算公式：(实际尺寸 × 焦距) / 像素尺寸
        distance = (self.armor_half_width * focal_length) / (pixel_distance * 0.5)
        return max(0.5, distance)  # 限制最小距离为0.5米

    def predict_ballistic_point(self, predicted_center, img_center, target_speed):
        """
        计算弹道预测点
        :param predicted_center: 卡尔曼预测的目标中心（图像坐标，(x,y)）
        :param img_center: 图像中心坐标（(x,y)）
        :param target_speed: 目标移动速度（由卡尔曼滤波器获取，(vx, vy)像素/秒）
        :return: 弹道修正后的预测点（图像坐标，(x,y)）
        """
        if predicted_center is None:
            self.last_ballistic_point = None
            return None

        # 1. 估算目标距离（米）
        distance = self._calculate_distance(predicted_center, img_center)

        # 2. 计算子弹飞行时间（秒）：距离 / 子弹速度
        flight_time = distance / self.bullet_speed if self.bullet_speed > 0 else 0.0

        # 3. 计算飞行时间内目标的移动量（像素）
        # 目标速度（pixel/s）× 飞行时间（s）= 预测移动像素
        delta_x = target_speed[0] * flight_time
        delta_y = target_speed[1] * flight_time

        # 4. 弹道预测点 = 卡尔曼预测中心 + 目标移动补偿
        ballistic_x = predicted_center[0] + delta_x
        ballistic_y = predicted_center[1] + delta_y
        ballistic_point = (ballistic_x, ballistic_y)

        self.last_ballistic_point = ballistic_point
        return ballistic_point

    def get_last_ballistic_point(self):
        """获取上一次弹道预测点"""
        return self.last_ballistic_point