import cv2
import numpy as np
from 视觉代码python.config import KF_PROCESS_NOISE, KF_MEASUREMENT_NOISE

class KalmanFilter:
    """卡尔曼滤波器：平滑并预测装甲板位置"""
    def __init__(self):
        self.kf = cv2.KalmanFilter(4, 2)  # 4状态(x,y,dx,dy)，2测量(x,y)
        self._init_params()
        self.initialized = False
        self.last_measurement = None
        self.last_prediction = None

    def _init_params(self):
        """初始化卡尔曼参数"""
        self.kf.transitionMatrix = np.array([
            [1, 0, 1, 0],
            [0, 1, 0, 1],
            [0, 0, 1, 0],
            [0, 0, 0, 1]
        ], np.float32)

        self.kf.measurementMatrix = np.array([
            [1, 0, 0, 0],
            [0, 1, 0, 0]
        ], np.float32)

        self.kf.processNoiseCov = np.eye(4, dtype=np.float32) * KF_PROCESS_NOISE
        self.kf.measurementNoiseCov = np.eye(2, dtype=np.float32) * KF_MEASUREMENT_NOISE
        self.kf.errorCovPost = np.eye(4, dtype=np.float32)

    def update(self, measurement):
        """用测量值更新滤波器"""
        measurement = np.array([[np.float32(measurement[0])], [np.float32(measurement[1])]])

        if not self.initialized:
            self.kf.statePost = np.array([
                [measurement[0, 0]],
                [measurement[1, 0]],
                [0.0],
                [0.0]
            ], np.float32)
            self.initialized = True
        else:
            self.kf.correct(measurement)

        self.last_measurement = (measurement[0, 0], measurement[1, 0])

    def predict(self):
        """预测下一帧位置"""
        if not self.initialized:
            return None

        prediction = self.kf.predict()
        self.last_prediction = (prediction[0, 0], prediction[1, 0])
        return self.last_prediction