"""
kalman_tracker.py — Single-object Kalman filter for 2-D ground-plane tracking.

State vector:  x = [px, py, vx, vy]   (position + velocity)
Motion model:  constant-velocity (CV) with process noise on acceleration

Predict step:
    F = | 1 0 dt  0 |      x_{k|k-1} = F @ x_{k-1|k-1}
        | 0 1  0 dt |      P_{k|k-1} = F @ P @ F^T + Q
        | 0 0  1  0 |
        | 0 0  0  1 |

Update step (position observation only):
    H = | 1 0 0 0 |      y   = z - H @ x
        | 0 1 0 0 |      S   = H @ P @ H^T + R
                          K   = P @ H^T @ inv(S)
                          x   = x + K @ y
                          P   = (I - K @ H) @ P   (Joseph form for stability)

Trajectory prediction: unrolls the CV model forward at fixed time steps.
"""

from __future__ import annotations

from typing import List, Tuple

import numpy as np


# ── Default noise matrices ────────────────────────────────────────────────────
# Process noise: models uncertainty in acceleration between frames
_Q_BASE = np.diag([0.02, 0.02, 0.8, 0.8]).astype(np.float64)

# Measurement noise: LIDAR centroid uncertainty (~0.15 m std)
_R = np.diag([0.025, 0.025]).astype(np.float64)   # 0.16 m sigma each axis

# Observation matrix
_H = np.array([[1, 0, 0, 0],
               [0, 1, 0, 0]], dtype=np.float64)
_I4 = np.eye(4, dtype=np.float64)


class KalmanTracker:
    """
    Kalman filter tracking one object.

    Parameters
    ----------
    x0, y0         : initial position (metres, odom frame)
    process_noise  : scalar multiplier on _Q_BASE
    """

    def __init__(self, x0: float, y0: float, process_noise: float = 1.0):
        self._x = np.array([x0, y0, 0.0, 0.0], dtype=np.float64)
        self._P = np.eye(4, dtype=np.float64) * 0.5
        self._Q = _Q_BASE * process_noise

    # ── Core filter ───────────────────────────────────────────────────────────

    def predict(self, dt: float) -> None:
        """Propagate state by dt seconds."""
        F = np.array([
            [1, 0, dt, 0],
            [0, 1,  0, dt],
            [0, 0,  1,  0],
            [0, 0,  0,  1],
        ], dtype=np.float64)
        self._x = F @ self._x
        self._P = F @ self._P @ F.T + self._Q

    def update(self, z: np.ndarray) -> None:
        """
        Incorporate a position measurement z = [x, y].
        Uses Joseph-form covariance update for numerical stability.
        """
        y = z.astype(np.float64) - _H @ self._x
        S = _H @ self._P @ _H.T + _R
        K = self._P @ _H.T @ np.linalg.inv(S)
        self._x = self._x + K @ y
        IKH = _I4 - K @ _H
        # Joseph form: (I-KH) P (I-KH)^T + K R K^T
        self._P = IKH @ self._P @ IKH.T + K @ _R @ K.T

    # ── Prediction horizon ────────────────────────────────────────────────────

    def predict_horizon(
        self,
        horizon_s: float = 2.5,
        step_s: float = 0.5,
    ) -> List[Tuple[float, float, float]]:
        """
        Return [(x, y, t), ...] at equally-spaced future times.

        Does NOT modify internal filter state.
        """
        predictions: List[Tuple[float, float, float]] = []
        state = self._x.copy()
        t = 0.0
        F_step = np.array([
            [1, 0, step_s,      0],
            [0, 1,      0, step_s],
            [0, 0,      1,      0],
            [0, 0,      0,      1],
        ], dtype=np.float64)
        while t < horizon_s - 1e-6:
            state = F_step @ state
            t += step_s
            predictions.append((float(state[0]), float(state[1]), t))
        return predictions

    # ── Properties ────────────────────────────────────────────────────────────

    @property
    def position(self) -> Tuple[float, float]:
        return float(self._x[0]), float(self._x[1])

    @property
    def velocity(self) -> Tuple[float, float]:
        return float(self._x[2]), float(self._x[3])

    @property
    def speed(self) -> float:
        return float(np.hypot(self._x[2], self._x[3]))

    @property
    def covariance_2x2(self) -> np.ndarray:
        """Position covariance (top-left 2×2 of P)."""
        return self._P[:2, :2].copy()

    @property
    def state(self) -> np.ndarray:
        return self._x.copy()
