"""
ByteTrack Object Tracker

Implementasi ByteTrack untuk multi-object tracking pada kendaraan.
Menggunakan Kalman Filter untuk prediksi state dan Hungarian Algorithm
untuk data association.

Referensi paper: ByteTrack: Multi-Object Tracking by Associating Every Detection Box
https://arxiv.org/abs/2110.06864
"""

import numpy as np
from scipy.optimize import linear_sum_assignment
from collections import deque


class KalmanFilter:
    """
    Kalman Filter sederhana untuk tracking bounding box.
    
    State vector: [cx, cy, w, h, vx, vy, vw, vh]
    - cx, cy: center x, y
    - w, h: width, height
    - vx, vy, vw, vh: velocity components
    """

    def __init__(self):
        # state transition matrix (8x8)
        self.F = np.eye(8)
        self.F[0, 4] = 1  # cx += vx
        self.F[1, 5] = 1  # cy += vy
        self.F[2, 6] = 1  # w += vw
        self.F[3, 7] = 1  # h += vh

        # observation matrix (4x8), kita observe [cx, cy, w, h]
        self.H = np.zeros((4, 8))
        self.H[0, 0] = 1
        self.H[1, 1] = 1
        self.H[2, 2] = 1
        self.H[3, 3] = 1

        # process noise
        self.Q = np.eye(8) * 1.0
        self.Q[4:, 4:] *= 0.01  # velocity noise lebih kecil

        # measurement noise
        self.R = np.eye(4) * 1.0

        # state dan covariance
        self.x = np.zeros(8)
        self.P = np.eye(8) * 10.0

    def init_state(self, bbox):
        """
        Inisialisasi state dari bounding box [x1, y1, x2, y2].
        """
        cx = (bbox[0] + bbox[2]) / 2
        cy = (bbox[1] + bbox[3]) / 2
        w = bbox[2] - bbox[0]
        h = bbox[3] - bbox[1]

        self.x = np.array([cx, cy, w, h, 0, 0, 0, 0], dtype=np.float64)
        self.P = np.eye(8) * 10.0

    def predict(self):
        """Predict step."""
        self.x = self.F @ self.x
        self.P = self.F @ self.P @ self.F.T + self.Q
        return self.x[:4]  # return predicted [cx, cy, w, h]

    def update(self, bbox):
        """
        Update step dengan measurement baru.
        
        Args:
            bbox: [x1, y1, x2, y2]
        """
        cx = (bbox[0] + bbox[2]) / 2
        cy = (bbox[1] + bbox[3]) / 2
        w = bbox[2] - bbox[0]
        h = bbox[3] - bbox[1]
        z = np.array([cx, cy, w, h])

        # innovation
        y = z - self.H @ self.x
        S = self.H @ self.P @ self.H.T + self.R
        K = self.P @ self.H.T @ np.linalg.inv(S)

        self.x = self.x + K @ y
        self.P = (np.eye(8) - K @ self.H) @ self.P

    def get_bbox(self):
        """Return current state sebagai [x1, y1, x2, y2]."""
        cx, cy, w, h = self.x[:4]
        x1 = cx - w / 2
        y1 = cy - h / 2
        x2 = cx + w / 2
        y2 = cy + h / 2
        return [x1, y1, x2, y2]

    def get_center(self):
        """Return center point (cx, cy)."""
        return self.x[0], self.x[1]


class Track:
    """
    Representasi satu tracked object.
    
    Menyimpan state, history, dan metadata untuk setiap kendaraan yang dilacak.
    """

    _next_id = 1  # class variable untuk auto-increment ID

    def __init__(self, bbox, class_id, class_name, confidence):
        self.track_id = Track._next_id
        Track._next_id += 1

        self.kf = KalmanFilter()
        self.kf.init_state(bbox)

        self.class_id = class_id
        self.class_name = class_name
        self.confidence = confidence

        self.hits = 1           # berapa kali di-match
        self.age = 0            # umur track (frame)
        self.time_since_update = 0  # frame sejak terakhir di-update

        # simpan history centroid untuk counting
        self.center_history = deque(maxlen=50)
        cx = (bbox[0] + bbox[2]) / 2
        cy = (bbox[1] + bbox[3]) / 2
        self.center_history.append((cx, cy))

        self.is_confirmed = False  # confirmed setelah beberapa hits

    def predict(self):
        """Predict posisi berikutnya."""
        self.kf.predict()
        self.age += 1
        self.time_since_update += 1

    def update(self, bbox, class_id, class_name, confidence):
        """Update track dengan deteksi baru."""
        self.kf.update(bbox)
        self.class_id = class_id
        self.class_name = class_name
        self.confidence = confidence
        self.hits += 1
        self.time_since_update = 0

        cx, cy = self.kf.get_center()
        self.center_history.append((cx, cy))

        # confirm track setelah 3 hits
        if self.hits >= 3:
            self.is_confirmed = True

    def get_bbox(self):
        """Return bounding box saat ini."""
        return self.kf.get_bbox()

    def get_center(self):
        """Return center point saat ini."""
        return self.kf.get_center()

    @classmethod
    def reset_id_counter(cls):
        """Reset ID counter, panggil di awal video baru."""
        cls._next_id = 1


def compute_iou(bbox1, bbox2):
    """
    Hitung Intersection over Union antara dua bounding box.
    
    Args:
        bbox1, bbox2: [x1, y1, x2, y2]
    
    Returns:
        float: IoU value
    """
    x1 = max(bbox1[0], bbox2[0])
    y1 = max(bbox1[1], bbox2[1])
    x2 = min(bbox1[2], bbox2[2])
    y2 = min(bbox1[3], bbox2[3])

    intersection = max(0, x2 - x1) * max(0, y2 - y1)

    area1 = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
    area2 = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])

    union = area1 + area2 - intersection

    if union <= 0:
        return 0.0

    return intersection / union


def compute_iou_matrix(bboxes1, bboxes2):
    """
    Hitung IoU matrix antara dua set bounding boxes.
    
    Returns:
        numpy array shape (len(bboxes1), len(bboxes2))
    """
    n = len(bboxes1)
    m = len(bboxes2)
    iou_matrix = np.zeros((n, m))

    for i in range(n):
        for j in range(m):
            iou_matrix[i, j] = compute_iou(bboxes1[i], bboxes2[j])

    return iou_matrix


class ByteTracker:
    """
    ByteTrack multi-object tracker.
    
    Implementasi berdasarkan paper ByteTrack yang menggunakan
    two-stage association untuk memanfaatkan deteksi low-confidence
    yang biasanya dibuang oleh tracker lain.
    
    Args:
        track_thresh: threshold untuk membedakan high/low confidence detections
        match_thresh: minimum IoU untuk matching
        track_buffer: jumlah frame sebelum track dihapus (kalau tidak di-update)
    """

    def __init__(self, track_thresh=0.5, match_thresh=0.3, track_buffer=30):
        self.track_thresh = track_thresh
        self.match_thresh = match_thresh
        self.track_buffer = track_buffer

        self.active_tracks = []     # track yang sedang aktif
        self.lost_tracks = []       # track yang hilang tapi belum dihapus

    def update(self, detections):
        """
        Update tracker dengan deteksi baru dari frame saat ini.
        
        Ini inti dari ByteTrack: two-stage association.
        Stage 1: match high-confidence detections dengan active tracks
        Stage 2: match low-confidence detections dengan unmatched tracks
        
        Args:
            detections: list of dict dari VehicleDetector.detect()
                masing-masing harus punya: bbox, confidence, class_id, class_name
        
        Returns:
            list of dict, setiap item berisi:
                - track_id: unique ID
                - bbox: [x1, y1, x2, y2]
                - class_id: int
                - class_name: str
                - confidence: float
                - center: (cx, cy)
        """
        # predict semua active tracks dulu
        for track in self.active_tracks:
            track.predict()
        for track in self.lost_tracks:
            track.predict()

        # pisahkan deteksi jadi high confidence dan low confidence
        high_dets = []
        low_dets = []

        for det in detections:
            if det["confidence"] >= self.track_thresh:
                high_dets.append(det)
            else:
                low_dets.append(det)

        # ============ Stage 1: match high-confidence dets dengan active tracks ============
        unmatched_tracks_idx = list(range(len(self.active_tracks)))
        unmatched_dets_idx = list(range(len(high_dets)))

        if len(self.active_tracks) > 0 and len(high_dets) > 0:
            track_bboxes = [t.get_bbox() for t in self.active_tracks]
            det_bboxes = [d["bbox"] for d in high_dets]

            iou_matrix = compute_iou_matrix(track_bboxes, det_bboxes)
            cost_matrix = 1.0 - iou_matrix  # karena hungarian minimize cost

            row_indices, col_indices = linear_sum_assignment(cost_matrix)

            matched_tracks = set()
            matched_dets = set()

            for row, col in zip(row_indices, col_indices):
                if iou_matrix[row, col] >= self.match_thresh:
                    # match berhasil
                    self.active_tracks[row].update(
                        high_dets[col]["bbox"],
                        high_dets[col]["class_id"],
                        high_dets[col]["class_name"],
                        high_dets[col]["confidence"]
                    )
                    matched_tracks.add(row)
                    matched_dets.add(col)

            unmatched_tracks_idx = [i for i in range(len(self.active_tracks)) if i not in matched_tracks]
            unmatched_dets_idx = [i for i in range(len(high_dets)) if i not in matched_dets]

        # ============ Stage 2: match low-confidence dets dengan unmatched tracks ============
        remaining_tracks = [self.active_tracks[i] for i in unmatched_tracks_idx]
        still_unmatched_tracks = list(range(len(remaining_tracks)))

        if len(remaining_tracks) > 0 and len(low_dets) > 0:
            track_bboxes = [t.get_bbox() for t in remaining_tracks]
            det_bboxes = [d["bbox"] for d in low_dets]

            iou_matrix = compute_iou_matrix(track_bboxes, det_bboxes)
            cost_matrix = 1.0 - iou_matrix

            row_indices, col_indices = linear_sum_assignment(cost_matrix)

            matched_in_stage2 = set()

            for row, col in zip(row_indices, col_indices):
                if iou_matrix[row, col] >= self.match_thresh:
                    remaining_tracks[row].update(
                        low_dets[col]["bbox"],
                        low_dets[col]["class_id"],
                        low_dets[col]["class_name"],
                        low_dets[col]["confidence"]
                    )
                    matched_in_stage2.add(row)

            still_unmatched_tracks = [i for i in range(len(remaining_tracks)) if i not in matched_in_stage2]

        # handle unmatched tracks -> pindahkan ke lost
        for idx in still_unmatched_tracks:
            track = remaining_tracks[idx]
            if track.time_since_update > self.track_buffer:
                continue  # buang, sudah terlalu lama hilang
            self.lost_tracks.append(track)

        # handle unmatched detections -> buat track baru
        for idx in unmatched_dets_idx:
            det = high_dets[idx]
            new_track = Track(
                det["bbox"],
                det["class_id"],
                det["class_name"],
                det["confidence"]
            )
            self.active_tracks.append(new_track)

        # coba match lost tracks juga dengan unmatched high-confidence dets
        # (ini versi simplified, di paper asli lebih kompleks)

        # bersihkan tracks yang sudah expired dari lost
        self.lost_tracks = [
            t for t in self.lost_tracks
            if t.time_since_update <= self.track_buffer
        ]

        # update active tracks: buang yang sudah lama tidak di-update
        self.active_tracks = [
            t for t in self.active_tracks
            if t.time_since_update <= self.track_buffer
        ]

        # gabungkan kembali lost tracks yang di-match ke active
        # (simplified: lost tracks tetap di list terpisah)

        # return hasil tracking
        output = []
        for track in self.active_tracks:
            if not track.is_confirmed:
                continue  # skip track yang belum confirmed

            bbox = track.get_bbox()
            cx, cy = track.get_center()

            output.append({
                "track_id": track.track_id,
                "bbox": bbox,
                "class_id": track.class_id,
                "class_name": track.class_name,
                "confidence": track.confidence,
                "center": (cx, cy),
                "center_history": list(track.center_history)
            })

        return output

    def reset(self):
        """Reset tracker untuk video baru."""
        self.active_tracks = []
        self.lost_tracks = []
        Track.reset_id_counter()