"""多目标追踪过滤器（TargetTracker）。

对应 C++ 参考实现 reference-c++/target_tracker/src/target_tracker.cpp
（规格：docs/视觉追踪技术路线.md §5 关联 / §6 生命周期与锁定 / §7 六步主流程）。

线程模型（§9.5）：update() 由单线程（推理主循环）调用；仅对 last_result
缓存加锁，last_result() 返回受保护拷贝。
"""

import math
import threading
from dataclasses import dataclass, field
from typing import Optional, Sequence

from .kalman_box_filter import BoxCXCYWH
from .tracklet import BoxXYXY, Tracklet, compute_iou, to_center_size


@dataclass
class Detection:
    """单个检测框（§2.2 输入元素：[x1,y1,x2,y2,conf,cls_id]）。"""

    x1: float = 0.0
    y1: float = 0.0
    x2: float = 0.0
    y2: float = 0.0
    conf: float = 0.0
    cls_id: int = 0


@dataclass
class FrameShape:
    """画面尺寸（height, width）：仅 Coast 兜底选举时用于取画面中心。"""

    height: int = 0
    width: int = 0


@dataclass
class TrackResult:
    """每帧追踪输出（§2.3 九字段）。"""

    tracked: bool = False
    primary_id: Optional[int] = None
    center: tuple[float, float] = (0.0, 0.0)
    box: Optional[BoxXYXY] = None
    confidence: float = 0.0
    is_predicted: bool = False
    lost_frames: int = 0
    n_active: int = 0
    raw: Optional[BoxXYXY] = None


@dataclass
class TrackerConfig:
    """追踪器配置参数（§8.1 默认值取自规格）。"""

    max_lost_frames: int = 8
    max_association_dist: float = 200.0
    min_hits: int = 3
    dist_weight: float = 0.5
    iou_weight: float = 0.5


@dataclass
class _AssociationResult:
    matches: list[tuple[int, int]] = field(default_factory=list)
    unmatched_track_ids: list[int] = field(default_factory=list)
    unmatched_det_indices: list[int] = field(default_factory=list)


class TargetTracker:
    def __init__(self, config: Optional[TrackerConfig] = None) -> None:
        self._config = config if config is not None else TrackerConfig()
        self._next_id = 0
        self._tracks: dict[int, Tracklet] = {}
        self._frame_count = 0
        self._locked_primary_id: Optional[int] = None
        self._last_result = TrackResult()
        self._last_result_mutex = threading.Lock()

    @property
    def config(self) -> TrackerConfig:
        return self._config

    def get_config(self) -> TrackerConfig:
        return self._config

    def update(self, detections: Sequence[Detection],
               frame_shape: Optional[FrameShape] = None) -> TrackResult:
        # §7 六步主流程
        self._frame_count += 1

        # 入口防御：过滤坐标或置信度含 NaN/Inf 的检测（非有限值一旦进入
        # 卡尔曼状态将永久污染轨道并传导到主目标输出）
        valid_detections = [
            det for det in detections
            if all(math.isfinite(v) for v in (det.x1, det.y1, det.x2, det.y2, det.conf))
        ]

        # 步骤 1：所有存活轨道先行预测（推进一帧）
        for track in self._tracks.values():
            track.predict()

        # 步骤 2：贪婪关联（无检测时全部轨道进入 Coast 计数）
        if valid_detections:
            assoc = self._associate_detections(valid_detections)
        else:
            assoc = _AssociationResult(unmatched_track_ids=list(self._tracks.keys()))

        # 步骤 3：用匹配到的检测更新轨道
        for track_id, det_idx in assoc.matches:
            det = valid_detections[det_idx]
            raw_box = BoxXYXY(det.x1, det.y1, det.x2, det.y2)
            self._tracks[track_id].update(raw_box, det.conf, det.cls_id)

        # 步骤 4：未匹配轨道计丢失，死亡者回收（锁定唯一解除途径之一）
        for track_id in assoc.unmatched_track_ids:
            track = self._tracks.get(track_id)
            if track is None:
                continue
            track.mark_miss()
            if track.is_dead:
                if self._locked_primary_id == track_id:
                    self._locked_primary_id = None
                del self._tracks[track_id]

        # 步骤 5：未匹配检测 → 创建新轨道（next_id 从 0 递增永不复用）
        for det_idx in assoc.unmatched_det_indices:
            det = valid_detections[det_idx]
            measurement = to_center_size(BoxXYXY(det.x1, det.y1, det.x2, det.y2))
            self._tracks[self._next_id] = Tracklet(
                self._next_id, measurement, det.conf, det.cls_id,
                self._config.max_lost_frames)
            self._next_id += 1

        # 步骤 6：选取/延续主目标，组装输出并缓存
        result = self._select_primary(frame_shape)
        with self._last_result_mutex:
            self._last_result = result
        return result

    def last_result(self) -> TrackResult:
        with self._last_result_mutex:
            return self._last_result

    def reset(self) -> None:
        # §2.1：清空所有轨道、ID 计数器、锁定状态与缓存结果（next_id 归零）
        self._tracks.clear()
        self._next_id = 0
        self._frame_count = 0
        self._locked_primary_id = None
        with self._last_result_mutex:
            self._last_result = TrackResult()

    def active_count(self) -> int:
        return len(self._tracks)

    def active_tracks(self) -> dict[int, Tracklet]:
        return dict(self._tracks)

    def _associate_detections(self, detections: Sequence[Detection]) -> _AssociationResult:
        result = _AssociationResult()

        # §5.3：按轨道 confidence 降序（平局按 ID 升序保证确定性）。
        # 非有限置信度按 −∞ 处理（排最后），避免 NaN 破坏全序。
        def sort_key(track_id: int) -> tuple[float, int]:
            raw = self._tracks[track_id].confidence
            conf = raw if math.isfinite(raw) else -math.inf
            return (-conf, track_id)

        sorted_ids = sorted(self._tracks.keys(), key=sort_key)
        det_assigned = [False] * len(detections)

        for track_id in sorted_ids:
            track = self._tracks[track_id]
            # §5.3：最优代价初始 1.0，仅 cost 严格小于才匹配（cost ≥ 1.0 判为不可匹配）
            best_cost = 1.0
            best_det = -1
            for d, det in enumerate(detections):
                if det_assigned[d]:
                    continue
                det_cx = (det.x1 + det.x2) / 2.0
                det_cy = (det.y1 + det.y2) / 2.0
                dist = math.hypot(track.last_center[0] - det_cx,
                                  track.last_center[1] - det_cy)
                norm_dist = dist / max(self._config.max_association_dist, 1e-3)
                iou = (compute_iou(track.last_box,
                                   BoxXYXY(det.x1, det.y1, det.x2, det.y2))
                       if track.last_box is not None else 0.0)
                cost = (self._config.dist_weight * norm_dist +
                        self._config.iou_weight * (1.0 - iou))
                if cost < best_cost:
                    best_cost = cost
                    best_det = d
            if best_det >= 0:
                result.matches.append((track_id, best_det))
                det_assigned[best_det] = True
            else:
                result.unmatched_track_ids.append(track_id)

        result.unmatched_det_indices = [
            d for d, assigned in enumerate(det_assigned) if not assigned
        ]
        return result

    def _select_primary(self, frame_shape: Optional[FrameShape]) -> TrackResult:
        result = TrackResult(n_active=len(self._tracks))

        # §6.3：轨道表为空 → 清除锁定，返回 tracked=false
        if not self._tracks:
            self._locked_primary_id = None
            return result

        # §6.3 ① 锁定延续：锁定轨道存活则直接输出（Coast 中照样输出，
        # 不被高置信度目标抢占）
        chosen: Optional[Tracklet] = None
        is_predicted = False
        if self._locked_primary_id is not None:
            track = self._tracks.get(self._locked_primary_id)
            if track is not None:
                chosen = track
                is_predicted = track.is_coasting

        # §6.3 ② 重新选举（无锁定 / 锁定目标已死亡被回收）
        if chosen is None:
            measured = [t for t in self._tracks.values() if not t.is_coasting]
            coasting = [t for t in self._tracks.values() if t.is_coasting]

            if measured:
                # confirmed ← measured 中 hits ≥ min_hits 者；为空则退回 measured
                # （min_hits 只影响候选池优先级，不阻止输出）
                pool = [t for t in measured if t.hits >= self._config.min_hits]
                if not pool:
                    pool = measured
                chosen = pool[0]
                for t in pool:
                    if t.confidence > chosen.confidence:
                        chosen = t
                is_predicted = False
            elif coasting:
                # 无实测轨道：选离画面中心最近者（缺省按 640×480 中心 (320,240)）
                center_x = frame_shape.width / 2.0 if frame_shape and frame_shape.width > 0 else 320.0
                center_y = frame_shape.height / 2.0 if frame_shape and frame_shape.height > 0 else 240.0
                best_dist_sq = math.inf
                for t in coasting:
                    dx = t.last_center[0] - center_x
                    dy = t.last_center[1] - center_y
                    dist_sq = dx * dx + dy * dy
                    if dist_sq < best_dist_sq:
                        best_dist_sq = dist_sq
                        chosen = t
                is_predicted = True
            else:
                # 理论上不可达（轨道非空则 measured/coasting 必居其一）；防御清锁
                self._locked_primary_id = None
                return result
            self._locked_primary_id = chosen.track_id  # 选举后锁定

        # 组装输出（§2.3 九字段）；raw 仅实测时有值，Coast 时为空
        result.tracked = True
        result.primary_id = chosen.track_id
        result.center = chosen.last_center
        result.box = chosen.last_box
        result.confidence = chosen.confidence
        result.is_predicted = is_predicted
        result.lost_frames = chosen.lost_count
        result.raw = None if is_predicted else chosen.last_raw_box
        return result