"""轨道（Tracklet）与基础几何工具。

对应 C++ 参考实现 reference-c++/target_tracker/include/target_tracker/tracklet.h
（规格：docs/视觉追踪技术路线.md §3.2 / §5.2 / §7）。
"""

from dataclasses import dataclass
from typing import Optional

from .kalman_box_filter import BoxCXCYWH, KalmanBoxFilter


@dataclass
class BoxXYXY:
    """轴对齐边界框（像素坐标，左上/右下角表示）。"""

    x1: float = 0.0
    y1: float = 0.0
    x2: float = 0.0
    y2: float = 0.0


def to_center_size(box: BoxXYXY) -> BoxCXCYWH:
    # §7 TRACKLET_UPDATE：cx=(x1+x2)/2；cy=(y1+y2)/2；w=x2−x1；h=y2−y1
    return BoxCXCYWH((box.x1 + box.x2) / 2.0, (box.y1 + box.y2) / 2.0,
                     box.x2 - box.x1, box.y2 - box.y1)


def to_xyxy(box: BoxCXCYWH) -> BoxXYXY:
    # §7 TRACKLET_PREDICT 的 center_size 逆变换
    return BoxXYXY(box.cx - box.w / 2.0, box.cy - box.h / 2.0,
                   box.cx + box.w / 2.0, box.cy + box.h / 2.0)


def compute_iou(a: BoxXYXY, b: BoxXYXY) -> float:
    """标准轴对齐 IoU（§5.2）：任一框为空、交集 ≤ 0 或并集 ≤ 0 时返回 0。"""
    area_a = (a.x2 - a.x1) * (a.y2 - a.y1)
    area_b = (b.x2 - b.x1) * (b.y2 - b.y1)
    if area_a <= 0.0 or area_b <= 0.0:
        return 0.0

    inter_w = min(a.x2, b.x2) - max(a.x1, b.x1)
    inter_h = min(a.y2, b.y2) - max(a.y1, b.y1)
    if inter_w <= 0.0 or inter_h <= 0.0:
        return 0.0

    inter_area = inter_w * inter_h
    union_area = area_a + area_b - inter_area
    if union_area <= 0.0:
        return 0.0
    return inter_area / union_area


class Tracklet:
    """单条轨道：生命周期计数 + 辅助状态 + 卡尔曼滤波器（§3.2）。

    构造即创建轨道（§7 步骤 5）：用首次观测初始化滤波器，age=1、hits=1、
    lost_count=0，last_center / last_box 立即由观测重建（创建当帧即可参与输出）。
    """

    def __init__(self, track_id: int, measurement: BoxCXCYWH, conf: float,
                 cls_id: int, max_lost_frames: int) -> None:
        self.track_id = track_id
        self.kf = KalmanBoxFilter()
        self.lost_count = 0
        self.max_lost = max_lost_frames
        self.confidence = conf
        self.cls_id = cls_id
        self.age = 1
        self.hits = 1
        self.last_center: tuple[float, float] = (measurement.cx, measurement.cy)
        self.last_box: Optional[BoxXYXY] = to_xyxy(measurement)
        # 创建即源于实测检测：同步记录原始检测框（Coast 期间保留不更新）
        self.last_raw_box: Optional[BoxXYXY] = self.last_box
        self.kf.initiate(measurement)

    @property
    def is_dead(self) -> bool:
        # §6.1 注记 2：lost_count > max_lost（严格大于）
        return self.lost_count > self.max_lost

    @property
    def is_coasting(self) -> bool:
        # §3.2：lost_count > 0 即处于纯预测状态
        return self.lost_count > 0

    def predict(self) -> None:
        # §7 TRACKLET_PREDICT：预测成功时更新 last_center / last_box 且 age+1
        pred = self.kf.predict()
        if pred is not None:
            self.last_center = (pred.cx, pred.cy)
            self.last_box = to_xyxy(pred)
            self.age += 1

    def update(self, box: BoxXYXY, conf: float, cls_id: int) -> None:
        # §7 TRACKLET_UPDATE：框 → [cx,cy,w,h] → kf.update
        self.kf.update(to_center_size(box))
        self.last_center = ((box.x1 + box.x2) / 2.0, (box.y1 + box.y2) / 2.0)
        self.last_box = box
        self.last_raw_box = box
        self.confidence = conf
        self.cls_id = cls_id
        self.lost_count = 0
        self.hits += 1
        self.age += 1

    def mark_miss(self) -> None:
        # §7 步骤 4：未匹配计丢失
        self.lost_count += 1
        self.age += 1