"""
kalman_tracker.py — 4-state linear Kalman filter for 2-D position+velocity tracking.

State vector: [x, y, vx, vy]
Observation:  [x_meas, y_meas]

Process model: constant velocity with Wiener process acceleration noise.
Tuned to handle EUC rider speeds (20–30 km/h ≈ 5.5–8.3 m/s) with fast
acceleration transients.

Pure module — no ROS2 dependency; fully unit-testable.
"""

import math
import numpy as np


class KalmanTracker:
    """
    4-state Kalman filter: state = [x, y, vx, vy].

    Parameters
    ----------
    process_noise   : acceleration noise standard deviation (m/s²).
                      Higher values allow the filter to track rapid velocity
                      changes (EUC acceleration events).  Default 3.0 m/s².
    meas_noise_uwb  : UWB position measurement noise std-dev (m). Default 0.20 m.
    meas_noise_cam  : Camera position measurement noise std-dev (m). Default 0.12 m.
    """

    def __init__(
        self,
        process_noise: float = 3.0,
        meas_noise_uwb: float = 0.20,
        meas_noise_cam: float = 0.12,
    ):
        self._q = float(process_noise)
        self._r_uwb = float(meas_noise_uwb)
        self._r_cam = float(meas_noise_cam)

        # State [x, y, vx, vy]
        self._x = np.zeros(4)

        # Covariance — large initial uncertainty (10 m², 10 (m/s)²)
        self._P = np.diag([10.0, 10.0, 10.0, 10.0])

        # Observation matrix: H * x = [x, y]
        self._H = np.array([[1.0, 0.0, 0.0, 0.0],
                             [0.0, 1.0, 0.0, 0.0]])

        self._initialized = False

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    @property
    def initialized(self) -> bool:
        return self._initialized

    def initialize(self, x: float, y: float) -> None:
        """Seed the filter at position (x, y) with zero velocity."""
        self._x = np.array([x, y, 0.0, 0.0])
        self._P = np.diag([1.0, 1.0, 5.0, 5.0])
        self._initialized = True

    def predict(self, dt: float) -> None:
        """
        Advance the filter state by dt seconds.

        Uses a discrete Wiener process acceleration model so that positional
        uncertainty grows as O(dt^4/4) and velocity uncertainty as O(dt^2).
        This lets the filter coast accurately through short signal outages
        while still being responsive to EUC velocity changes.
        """
        if dt <= 0.0:
            return

        F = np.array([[1.0, 0.0,  dt, 0.0],
                      [0.0, 1.0, 0.0,  dt],
                      [0.0, 0.0, 1.0, 0.0],
                      [0.0, 0.0, 0.0, 1.0]])

        q = self._q
        dt2 = dt * dt
        dt3 = dt2 * dt
        dt4 = dt3 * dt
        Q = (q * q) * np.array([
            [dt4 / 4.0, 0.0,       dt3 / 2.0, 0.0      ],
            [0.0,       dt4 / 4.0, 0.0,       dt3 / 2.0],
            [dt3 / 2.0, 0.0,       dt2,       0.0      ],
            [0.0,       dt3 / 2.0, 0.0,       dt2      ],
        ])

        self._x = F @ self._x
        self._P = F @ self._P @ F.T + Q

    def update(self, x_meas: float, y_meas: float, source: str = "uwb") -> None:
        """
        Apply a position measurement (x_meas, y_meas).

        source : "uwb" or "camera" — selects the appropriate noise covariance.
        """
        r = self._r_uwb if source == "uwb" else self._r_cam
        R = np.diag([r * r, r * r])

        z = np.array([x_meas, y_meas])
        innov = z - self._H @ self._x                      # innovation
        S = self._H @ self._P @ self._H.T + R              # innovation covariance
        K = self._P @ self._H.T @ np.linalg.inv(S)        # Kalman gain
        self._x = self._x + K @ innov
        self._P = (np.eye(4) - K @ self._H) @ self._P

    # ------------------------------------------------------------------
    # State accessors
    # ------------------------------------------------------------------

    @property
    def position(self) -> tuple:
        """Current filtered position (x, y) in metres."""
        return float(self._x[0]), float(self._x[1])

    @property
    def velocity(self) -> tuple:
        """Current filtered velocity (vx, vy) in m/s."""
        return float(self._x[2]), float(self._x[3])

    def position_uncertainty_m(self) -> float:
        """RMS positional uncertainty (m) from diagonal of covariance."""
        return float(math.sqrt((self._P[0, 0] + self._P[1, 1]) / 2.0))

    def covariance_copy(self) -> np.ndarray:
        """Return a copy of the current 4×4 covariance matrix."""
        return self._P.copy()
