from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from scipy.optimize import linear_sum_assignment


def _box_to_state(box: np.ndarray) -> np.ndarray:
    x0, y0, x1, y1 = box
    width, height = x1 - x0, y1 - y0
    cx, cy = x0 + width / 2.0, y0 + height / 2.0
    scale = width * height
    aspect = width / max(height, 1e-6)
    return np.array([cx, cy, scale, aspect], dtype=np.float64)


def _state_to_box(state: np.ndarray) -> np.ndarray:
    cx, cy, scale, aspect = state[:4]
    scale = max(scale, 1e-6)
    width = np.sqrt(scale * aspect)
    height = scale / max(width, 1e-6)
    return np.array(
        [cx - width / 2.0, cy - height / 2.0, cx + width / 2.0, cy + height / 2.0],
        dtype=np.float64,
    )


def iou_matrix(boxes_a: np.ndarray, boxes_b: np.ndarray) -> np.ndarray:
    if boxes_a.shape[0] == 0 or boxes_b.shape[0] == 0:
        return np.zeros((boxes_a.shape[0], boxes_b.shape[0]), dtype=np.float64)
    area_a = (boxes_a[:, 2] - boxes_a[:, 0]) * (boxes_a[:, 3] - boxes_a[:, 1])
    area_b = (boxes_b[:, 2] - boxes_b[:, 0]) * (boxes_b[:, 3] - boxes_b[:, 1])
    top_left = np.maximum(boxes_a[:, None, :2], boxes_b[None, :, :2])
    bottom_right = np.minimum(boxes_a[:, None, 2:], boxes_b[None, :, 2:])
    width_height = np.clip(bottom_right - top_left, 0, None)
    intersection = width_height[..., 0] * width_height[..., 1]
    union = area_a[:, None] + area_b[None, :] - intersection
    return intersection / np.clip(union, 1e-9, None)


class _KalmanBoxTracker:
    """Constant-velocity Kalman filter over (cx, cy, scale, aspect) — the classic SORT state.

    Aspect ratio is treated as constant (no velocity term for it), matching Bewley et al. 2016.
    """

    _next_id = 1

    def __init__(self, box: np.ndarray, label: int, score: float) -> None:
        self.state = np.zeros(7, dtype=np.float64)
        self.state[:4] = _box_to_state(box)
        self.covariance = np.eye(7) * 10.0
        self.covariance[4:, 4:] *= 1000.0

        self._transition = np.eye(7)
        for i in range(3):
            self._transition[i, i + 4] = 1.0
        self._observation = np.zeros((4, 7))
        self._observation[:4, :4] = np.eye(4)

        self._process_noise = np.eye(7) * 1.0
        self._process_noise[4:, 4:] *= 0.01
        self._measurement_noise = np.eye(4) * 1.0

        self.id = _KalmanBoxTracker._next_id
        _KalmanBoxTracker._next_id += 1
        self.label = label
        self.score = score
        self.hits = 1
        self.age = 0
        self.time_since_update = 0

    def predict(self) -> np.ndarray:
        self.state = self._transition @ self.state
        self.covariance = self._transition @ self.covariance @ self._transition.T + self._process_noise
        self.age += 1
        self.time_since_update += 1
        state = self.state.copy()
        state[2] = max(state[2], 1e-6)
        return _state_to_box(state)

    def update(self, box: np.ndarray, label: int, score: float) -> None:
        measurement = _box_to_state(box)
        innovation = measurement - self._observation @ self.state
        innovation_cov = self._observation @ self.covariance @ self._observation.T + self._measurement_noise
        kalman_gain = self.covariance @ self._observation.T @ np.linalg.inv(innovation_cov)
        self.state = self.state + kalman_gain @ innovation
        self.covariance = (np.eye(7) - kalman_gain @ self._observation) @ self.covariance
        self.label = label
        self.score = score
        self.hits += 1
        self.time_since_update = 0

    def current_box(self) -> np.ndarray:
        return _state_to_box(self.state)


@dataclass
class Track:
    id: int
    box: tuple[float, float, float, float]
    label: int
    score: float
    hits: int
    age: int


class SortTracker:
    """Minimal SORT-style tracker: Kalman motion prediction + IoU/Hungarian association.

    This sits *outside* the model as a post-processing layer over independent
    per-frame detections — ObjectModel-v1 itself has no temporal component.
    Gives detections a persistent id and lets a track survive a few frames of
    missed detection (occlusion, a confidence dip) via `max_age`.
    """

    def __init__(self, iou_threshold: float = 0.3, max_age: int = 5, min_hits: int = 3) -> None:
        self.iou_threshold = iou_threshold
        self.max_age = max_age
        self.min_hits = min_hits
        self._trackers: list[_KalmanBoxTracker] = []

    def update(self, boxes: np.ndarray, labels: np.ndarray, scores: np.ndarray) -> list[Track]:
        """Advance one frame. `boxes` is [N, 4] xyxy, `labels`/`scores` are [N]."""
        predicted = (
            np.array([tracker.predict() for tracker in self._trackers])
            if self._trackers
            else np.zeros((0, 4))
        )

        matches, unmatched_detections, _ = self._associate(predicted, boxes)

        for det_idx, trk_idx in matches:
            self._trackers[trk_idx].update(boxes[det_idx], int(labels[det_idx]), float(scores[det_idx]))

        for det_idx in unmatched_detections:
            self._trackers.append(
                _KalmanBoxTracker(boxes[det_idx], int(labels[det_idx]), float(scores[det_idx]))
            )

        self._trackers = [t for t in self._trackers if t.time_since_update <= self.max_age]

        results = []
        for tracker in self._trackers:
            confirmed = tracker.hits >= self.min_hits or tracker.age <= self.min_hits
            if tracker.time_since_update == 0 and confirmed:
                x0, y0, x1, y1 = tracker.current_box()
                results.append(
                    Track(
                        id=tracker.id,
                        box=(x0, y0, x1, y1),
                        label=tracker.label,
                        score=tracker.score,
                        hits=tracker.hits,
                        age=tracker.age,
                    )
                )
        return results

    def _associate(
        self, predicted: np.ndarray, detections: np.ndarray
    ) -> tuple[list[tuple[int, int]], list[int], list[int]]:
        if predicted.shape[0] == 0 or detections.shape[0] == 0:
            return [], list(range(detections.shape[0])), list(range(predicted.shape[0]))

        iou = iou_matrix(detections, predicted)
        row_idx, col_idx = linear_sum_assignment(1.0 - iou)

        matches: list[tuple[int, int]] = []
        matched_detections: set[int] = set()
        matched_trackers: set[int] = set()
        for row, col in zip(row_idx, col_idx, strict=True):
            if iou[row, col] >= self.iou_threshold:
                matches.append((row, col))
                matched_detections.add(row)
                matched_trackers.add(col)

        unmatched_detections = [i for i in range(detections.shape[0]) if i not in matched_detections]
        unmatched_trackers = [i for i in range(predicted.shape[0]) if i not in matched_trackers]
        return matches, unmatched_detections, unmatched_trackers