"""
_obstacle_velocity.py — Dynamic obstacle velocity estimation via Kalman
filtering of LIDAR cluster centroids (no ROS2 deps).

Algorithm
---------
Each detected cluster centroid is tracked by a constant-velocity 2-D Kalman
filter:

  State    : [x, y, vx, vy]^T
  Obs      : [x, y]^T  (centroid position)

  Predict  : x_pred = F(dt) @ x;   P_pred = F @ P @ F^T + Q(dt)
  Update   : S = H @ P_pred @ H^T + R
             K = P_pred @ H^T @ inv(S)
             x = x_pred + K @ (z - H @ x_pred)
             P = (I - K @ H) @ P_pred

Process noise Q uses the white-noise-acceleration approximation:
  Q = diag([q_pos·dt², q_pos·dt², q_vel·dt, q_vel·dt])

Data association uses a greedy nearest-centroid approach: scan all
(track, cluster) pairs in order of Euclidean distance and accept matches
below `max_association_dist_m`, one-to-one.

Track lifecycle
---------------
  created   : first appearance of a cluster with no nearby track
  confident : after `n_init_frames` consecutive updates (confidence=1.0)
  coasting  : cluster not matched for up to `max_coasting_frames` calls
  deleted   : coasting count exceeds `max_coasting_frames`

Public API
----------
  KalmanTrack                           one tracked obstacle
    .predict(dt)                        advance state by dt seconds
    .update(centroid, width, depth, n)  correct with new measurement
    .position  → np.ndarray (2,)
    .velocity  → np.ndarray (2,)
    .speed     → float (m/s)
    .confidence → float 0–1
    .alive     → bool

  associate(positions, centroids, max_dist)
                                        → (matches, unmatched_t, unmatched_c)

  ObstacleTracker
    .update(centroids, timestamp, widths, depths, point_counts)
                                        → List[KalmanTrack]
    .tracks                             → Dict[int, KalmanTrack]
"""

from __future__ import annotations

from typing import Dict, List, Optional, Tuple

import numpy as np


# ── Kalman filter constants ───────────────────────────────────────────────────

_H = np.array([[1.0, 0.0, 0.0, 0.0],
               [0.0, 1.0, 0.0, 0.0]], dtype=np.float64)

_I4 = np.eye(4, dtype=np.float64)


# ── KalmanTrack ───────────────────────────────────────────────────────────────

class KalmanTrack:
    """
    Constant-velocity 2-D Kalman filter for one LIDAR cluster centroid.

    Parameters
    ----------
    track_id      : unique integer ID
    centroid      : initial (x, y) position in metres
    q_pos         : position process noise density  (m²/s²)
    q_vel         : velocity process noise density  (m²/s³)
    r_pos         : measurement noise std-dev       (metres)
    n_init_frames : consecutive updates needed for confidence=1
    max_coasting  : missed updates before track is deleted
    """

    def __init__(
        self,
        track_id:      int,
        centroid:      np.ndarray,
        q_pos:         float = 0.05,
        q_vel:         float = 0.50,
        r_pos:         float = 0.10,
        n_init_frames: int   = 3,
        max_coasting:  int   = 5,
    ) -> None:
        self._id          = track_id
        self._n_init      = max(n_init_frames, 1)
        self._max_coast   = max_coasting
        self._q_pos       = float(q_pos)
        self._q_vel       = float(q_vel)
        self._R           = np.diag([r_pos ** 2, r_pos ** 2])
        self._age         = 0       # successful updates
        self._coasting    = 0       # consecutive missed updates
        self._alive       = True

        # State: [x, y, vx, vy]^T — initial velocity is zero
        self._x = np.array(
            [float(centroid[0]), float(centroid[1]), 0.0, 0.0],
            dtype=np.float64,
        )
        # High initial velocity uncertainty, low position uncertainty
        self._P = np.diag([r_pos ** 2, r_pos ** 2, 10.0, 10.0])

        # Last-known cluster metadata (stored, not filtered)
        self.last_width:       float = 0.0
        self.last_depth:       float = 0.0
        self.last_point_count: int   = 0

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

    @property
    def track_id(self) -> int:
        return self._id

    @property
    def position(self) -> np.ndarray:
        return self._x[:2].copy()

    @property
    def velocity(self) -> np.ndarray:
        return self._x[2:].copy()

    @property
    def speed(self) -> float:
        return float(np.linalg.norm(self._x[2:]))

    @property
    def confidence(self) -> float:
        return min(1.0, self._age / self._n_init)

    @property
    def alive(self) -> bool:
        return self._alive

    @property
    def age(self) -> int:
        return self._age

    @property
    def coasting(self) -> int:
        return self._coasting

    # ── Kalman steps ──────────────────────────────────────────────────────────

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

        dt must be ≥ 0; negative values are clamped to 0.
        Increments the coasting counter; marks the track dead when the
        counter exceeds max_coasting_frames.
        """
        dt = max(0.0, float(dt))

        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],
        ], dtype=np.float64)

        Q = np.diag([
            self._q_pos * dt * dt,
            self._q_pos * dt * dt,
            self._q_vel * dt,
            self._q_vel * dt,
        ])

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

        self._coasting += 1
        if self._coasting > self._max_coast:
            self._alive = False

    def update(
        self,
        centroid:    np.ndarray,
        width:       float = 0.0,
        depth:       float = 0.0,
        point_count: int   = 0,
    ) -> None:
        """
        Correct state with a new centroid measurement.

        Resets the coasting counter and increments the age counter.
        """
        z = np.asarray(centroid, dtype=np.float64)[:2]

        S = _H @ self._P @ _H.T + self._R          # innovation covariance
        K = self._P @ _H.T @ np.linalg.inv(S)      # Kalman gain

        self._x = self._x + K @ (z - _H @ self._x)
        self._P = (_I4 - K @ _H) @ self._P

        self._age       += 1
        self._coasting   = 0

        self.last_width       = float(width)
        self.last_depth       = float(depth)
        self.last_point_count = int(point_count)


# ── Data association ──────────────────────────────────────────────────────────

def associate(
    positions:  np.ndarray,   # (N, 2) predicted track positions
    centroids:  np.ndarray,   # (M, 2) new cluster centroids
    max_dist:   float,
) -> Tuple[List[Tuple[int, int]], List[int], List[int]]:
    """
    Greedy nearest-centroid data association.

    Iteratively matches the closest (track, cluster) pair, one-to-one, as
    long as the distance is strictly less than `max_dist`.

    Returns
    -------
    matches           : list of (track_idx, cluster_idx) pairs
    unmatched_tracks  : track indices with no assigned cluster
    unmatched_clusters: cluster indices with no assigned track
    """
    N = len(positions)
    M = len(centroids)

    if N == 0 or M == 0:
        return [], list(range(N)), list(range(M))

    # (N, M) pairwise Euclidean distance matrix
    cost = np.linalg.norm(
        positions[:, None, :] - centroids[None, :, :], axis=2
    ).astype(np.float64)   # (N, M)

    matched_t: set = set()
    matched_c: set = set()
    matches: List[Tuple[int, int]] = []

    for _ in range(min(N, M)):
        if cost.min() >= max_dist:
            break
        t_idx, c_idx = np.unravel_index(int(cost.argmin()), cost.shape)
        matches.append((int(t_idx), int(c_idx)))
        matched_t.add(t_idx)
        matched_c.add(c_idx)
        cost[t_idx, :] = np.inf
        cost[:, c_idx] = np.inf

    unmatched_t = [i for i in range(N) if i not in matched_t]
    unmatched_c = [i for i in range(M) if i not in matched_c]
    return matches, unmatched_t, unmatched_c


# ── ObstacleTracker ───────────────────────────────────────────────────────────

class ObstacleTracker:
    """
    Multi-obstacle Kalman tracker operating on LIDAR cluster centroids.

    Call `update()` once per LIDAR scan with the list of cluster centroids
    (and optional bbox / point_count metadata).  Returns the current list
    of alive KalmanTrack objects.
    """

    def __init__(
        self,
        max_association_dist_m: float = 0.50,
        max_coasting_frames:    int   = 5,
        n_init_frames:          int   = 3,
        q_pos:                  float = 0.05,
        q_vel:                  float = 0.50,
        r_pos:                  float = 0.10,
    ) -> None:
        self._max_dist   = float(max_association_dist_m)
        self._max_coast  = int(max_coasting_frames)
        self._n_init     = int(n_init_frames)
        self._q_pos      = float(q_pos)
        self._q_vel      = float(q_vel)
        self._r_pos      = float(r_pos)

        self._tracks:   Dict[int, KalmanTrack] = {}
        self._next_id:  int   = 1
        self._last_t:   Optional[float] = None

    @property
    def tracks(self) -> Dict[int, KalmanTrack]:
        return self._tracks

    def update(
        self,
        centroids:    List[np.ndarray],
        timestamp:    float,
        widths:       Optional[List[float]] = None,
        depths:       Optional[List[float]] = None,
        point_counts: Optional[List[int]]   = None,
    ) -> List[KalmanTrack]:
        """
        Process one frame of cluster centroids.

        Parameters
        ----------
        centroids    : list of (2,) arrays — cluster centroid positions
        timestamp    : wall-clock time in seconds (time.time())
        widths       : optional per-cluster bbox widths (metres)
        depths       : optional per-cluster bbox depths (metres)
        point_counts : optional per-cluster LIDAR point counts

        Returns
        -------
        List of all alive KalmanTrack objects after this update.
        """
        dt = max(0.0, timestamp - self._last_t) if self._last_t is not None else 0.0
        self._last_t = timestamp

        # 1. Predict all existing tracks forward by dt
        for track in self._tracks.values():
            track.predict(dt)

        # 2. Build arrays for association (alive tracks only)
        alive = [t for t in self._tracks.values() if t.alive]
        if alive:
            pred_pos = np.array([t.position for t in alive])   # (N, 2)
        else:
            pred_pos = np.empty((0, 2), dtype=np.float64)

        if centroids:
            cent_arr = np.array([c[:2] for c in centroids], dtype=np.float64)  # (M, 2)
        else:
            cent_arr = np.empty((0, 2), dtype=np.float64)

        # 3. Associate
        matches, _, unmatched_c = associate(pred_pos, cent_arr, self._max_dist)

        # 4. Update matched tracks
        for ti, ci in matches:
            track = alive[ti]
            track.update(
                cent_arr[ci],
                width       = widths[ci]       if widths       else 0.0,
                depth       = depths[ci]       if depths       else 0.0,
                point_count = point_counts[ci] if point_counts else 0,
            )

        # 5. Spawn new tracks for unmatched clusters
        for ci in unmatched_c:
            track = KalmanTrack(
                track_id      = self._next_id,
                centroid      = cent_arr[ci],
                q_pos         = self._q_pos,
                q_vel         = self._q_vel,
                r_pos         = self._r_pos,
                n_init_frames = self._n_init,
                max_coasting  = self._max_coast,
            )
            track.update(
                cent_arr[ci],
                width       = widths[ci]       if widths       else 0.0,
                depth       = depths[ci]       if depths       else 0.0,
                point_count = point_counts[ci] if point_counts else 0,
            )
            self._tracks[self._next_id] = track
            self._next_id += 1

        # 6. Prune dead tracks
        self._tracks = {tid: t for tid, t in self._tracks.items() if t.alive}

        return list(self._tracks.values())
