import numpy as np
from filterpy.kalman import KalmanFilter

class BallTracker:
    def __init__(self):
        self.kf = KalmanFilter(dim_x=4, dim_z=2)
        self.kf.F = np.array([[1,0,1,0], [0,1,0,1], [0,0,1,0], [0,0,0,1]])
        self.kf.H = np.array([[1,0,0,0], [0,1,0,0]])
        self.kf.P *= 10.
        self.kf.R *= 0.5
        self.kf.Q *= 0.01
        self.initialized = False

    def update(self, measurement):
        if not self.initialized:
            self.kf.x[:2] = np.array([[measurement[0]], [measurement[1]]])
            self.initialized = True
        self.kf.predict()
        self.kf.update(measurement)
        return self.kf.x[:2].flatten()