"""YOLO obstacle detection with ByteTrack tracking and state management."""
from config import YOLO_MODEL, YOLO_CONF, OBSTACLE_CLASSES

_yolo = None


def _load():
    global _yolo
    if _yolo is None:
        from ultralytics import YOLO
        _yolo = YOLO(YOLO_MODEL)
    return _yolo


def detect(frame, track=False):
    """Detect obstacles. If track=True, uses ByteTrack for persistent IDs."""
    yolo = _load()
    h, w = frame.shape[:2]

    if track:
        results = yolo.track(frame, conf=YOLO_CONF, verbose=False, persist=True)[0]
    else:
        results = yolo(frame, conf=YOLO_CONF, verbose=False)[0]

    obstacles = []
    for box in results.boxes:
        cls_id = int(box.cls[0])
        if cls_id not in OBSTACLE_CLASSES:
            continue
        x1, y1, x2, y2 = map(int, box.xyxy[0])

        # Filter out full-frame false positives (box covers >50% of frame area)
        box_area = (x2 - x1) * (y2 - y1)
        if box_area > 0.5 * h * w:
            continue

        proximity = y2 / h
        cx = (x1 + x2) / 2
        direction = "LEFT" if cx < w * 0.33 else "RIGHT" if cx > w * 0.66 else "CENTER"

        ob = {
            'label': OBSTACLE_CLASSES[cls_id],
            'conf': float(box.conf[0]),
            'box': (x1, y1, x2, y2),
            'center': ((x1 + x2) // 2, (y1 + y2) // 2),
            'proximity': proximity,
            'dist': "NEAR" if proximity > 0.7 else "MID" if proximity > 0.4 else "FAR",
            'direction': direction,
        }
        if track and box.id is not None:
            ob['track_id'] = int(box.id[0])
        obstacles.append(ob)

    return sorted(obstacles, key=lambda o: -o['proximity'])


class ObstacleTracker:
    """Tracks obstacles across frames, detects new/closing objects."""

    def __init__(self):
        self._prev = {}       # track_id -> previous obstacle dict
        self._new_ids = set()  # track_ids that appeared this frame
        self._lost_ids = set() # track_ids that disappeared this frame

    def update(self, obstacles):
        """Update tracker state. Call once per frame after detect(track=True).

        Returns (new_obstacles, closing_obstacles, lost_ids).
          - new_obstacles: obstacles with track_ids not seen before
          - closing_obstacles: obstacles whose proximity increased significantly
          - lost_ids: track_ids from previous frame no longer present
        """
        current = {}
        new_obs = []
        closing_obs = []

        for ob in obstacles:
            tid = ob.get('track_id')
            if tid is None:
                continue
            current[tid] = ob

            if tid not in self._prev:
                new_obs.append(ob)
            else:
                # Check if closing (proximity increasing = getting nearer)
                prev_prox = self._prev[tid]['proximity']
                delta = ob['proximity'] - prev_prox
                if delta > 0.05:  # moved noticeably closer
                    ob['closing_rate'] = round(delta, 3)
                    closing_obs.append(ob)

        self._lost_ids = set(self._prev.keys()) - set(current.keys())
        self._new_ids = set(current.keys()) - set(self._prev.keys())
        self._prev = current

        return new_obs, closing_obs, self._lost_ids

    @property
    def active_count(self):
        return len(self._prev)