"""
tracker_manager.py — Multi-object tracker with Hungarian data association.

Track lifecycle:
  TENTATIVE → confirmed after `confirm_frames` consecutive hits
  CONFIRMED → normal tracked state
  LOST      → missed for 1..max_missed frames (still predicts, not published)
  DEAD      → missed > max_missed → removed

Association:
  Uses scipy.optimize.linear_sum_assignment (Hungarian algorithm) on a cost
  matrix of Euclidean distances between predicted track positions and new
  detections.  Assignments with cost > assoc_dist_m are rejected.

Up to `max_tracks` simultaneous live tracks (tentative + confirmed).
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import IntEnum
from typing import Dict, List, Optional, Tuple

import numpy as np
from scipy.optimize import linear_sum_assignment

from .kalman_tracker import KalmanTracker
from .object_detector import Detection


class TrackState(IntEnum):
    TENTATIVE = 0
    CONFIRMED = 1
    LOST      = 2


@dataclass
class Track:
    track_id:   int
    kalman:     KalmanTracker
    state:      TrackState = TrackState.TENTATIVE
    hits:       int = 1
    age:        int = 1      # frames since creation
    missed:     int = 0      # consecutive missed frames


class TrackerManager:
    """
    Manages a pool of Kalman tracks.

    Parameters
    ----------
    max_tracks      : hard cap on simultaneously alive tracks
    confirm_frames  : hits needed before a track is CONFIRMED
    max_missed      : consecutive missed frames before a track is DEAD
    assoc_dist_m    : max allowed distance (m) for a valid assignment
    horizon_s       : prediction horizon for trajectory output (seconds)
    pred_step_s     : time step between predicted waypoints
    process_noise   : KalmanTracker process-noise multiplier
    """

    def __init__(
        self,
        max_tracks:     int   = 20,
        confirm_frames: int   = 3,
        max_missed:     int   = 6,
        assoc_dist_m:   float = 1.5,
        horizon_s:      float = 2.5,
        pred_step_s:    float = 0.5,
        process_noise:  float = 1.0,
    ):
        self._max_tracks     = max_tracks
        self._confirm_frames = confirm_frames
        self._max_missed     = max_missed
        self._assoc_dist     = assoc_dist_m
        self._horizon_s      = horizon_s
        self._pred_step      = pred_step_s
        self._proc_noise     = process_noise

        self._tracks:  Dict[int, Track] = {}
        self._next_id: int = 1

    # ── Public API ────────────────────────────────────────────────────────────

    def update(self, detections: List[Detection], dt: float) -> List[Track]:
        """
        Process one frame of detections.

        1. Predict all tracks by dt.
        2. Hungarian assignment of predictions → detections.
        3. Update matched tracks; mark unmatched tracks as LOST.
        4. Promote tracks crossing `confirm_frames`.
        5. Create new tracks for unmatched detections (if room).
        6. Remove DEAD tracks.

        Returns confirmed tracks only.
        """
        # 1. Predict
        for tr in self._tracks.values():
            tr.kalman.predict(dt)
            tr.age += 1

        # 2. Assign
        matched, unmatched_tracks, unmatched_dets = self._assign(detections)

        # 3a. Update matched
        for tid, did in matched:
            tr = self._tracks[tid]
            det = detections[did]
            tr.kalman.update(np.array([det.x, det.y]))
            tr.hits   += 1
            tr.missed  = 0
            if tr.state == TrackState.LOST:
                tr.state = TrackState.CONFIRMED
            elif tr.state == TrackState.TENTATIVE and tr.hits >= self._confirm_frames:
                tr.state = TrackState.CONFIRMED

        # 3b. Unmatched tracks
        for tid in unmatched_tracks:
            tr = self._tracks[tid]
            tr.missed += 1
            if tr.missed > 1:
                tr.state = TrackState.LOST

        # 4. New tracks for unmatched detections
        live = sum(1 for t in self._tracks.values() if t.state != TrackState.LOST
                   or t.missed <= self._max_missed)
        for did in unmatched_dets:
            if live >= self._max_tracks:
                break
            det = detections[did]
            init_state = (TrackState.CONFIRMED
                          if self._confirm_frames <= 1
                          else TrackState.TENTATIVE)
            new_tr = Track(
                track_id=self._next_id,
                kalman=KalmanTracker(det.x, det.y, self._proc_noise),
                state=init_state,
            )
            self._tracks[self._next_id] = new_tr
            self._next_id += 1
            live += 1

        # 5. Prune dead
        dead = [tid for tid, t in self._tracks.items() if t.missed > self._max_missed]
        for tid in dead:
            del self._tracks[tid]

        return [t for t in self._tracks.values() if t.state == TrackState.CONFIRMED]

    @property
    def all_tracks(self) -> List[Track]:
        return list(self._tracks.values())

    @property
    def tentative_count(self) -> int:
        return sum(1 for t in self._tracks.values()
                   if t.state == TrackState.TENTATIVE)

    def reset(self) -> None:
        self._tracks.clear()
        self._next_id = 1

    # ── Hungarian assignment ──────────────────────────────────────────────────

    def _assign(
        self,
        detections: List[Detection],
    ) -> Tuple[List[Tuple[int, int]], List[int], List[int]]:
        """
        Returns:
          matched        — list of (track_id, det_index)
          unmatched_tids — track IDs with no detection assigned
          unmatched_dids — detection indices with no track assigned
        """
        track_ids = list(self._tracks.keys())

        if not track_ids or not detections:
            return [], track_ids, list(range(len(detections)))

        # Build cost matrix: rows=tracks, cols=detections
        cost = np.full((len(track_ids), len(detections)), fill_value=np.inf)
        for r, tid in enumerate(track_ids):
            tx, ty = self._tracks[tid].kalman.position
            for c, det in enumerate(detections):
                cost[r, c] = np.hypot(tx - det.x, ty - det.y)

        row_ind, col_ind = linear_sum_assignment(cost)

        matched:          List[Tuple[int, int]] = []
        matched_track_idx: set = set()
        matched_det_idx:   set = set()

        for r, c in zip(row_ind, col_ind):
            if cost[r, c] > self._assoc_dist:
                continue
            matched.append((track_ids[r], c))
            matched_track_idx.add(r)
            matched_det_idx.add(c)

        unmatched_tids = [track_ids[r] for r in range(len(track_ids))
                          if r not in matched_track_idx]
        unmatched_dids = [c for c in range(len(detections))
                          if c not in matched_det_idx]

        return matched, unmatched_tids, unmatched_dids
