import numpy as np
from scipy.optimize import linear_sum_assignment

def read_tracking_file(file_path):
    """ Read tracking results or ground truth file and return a dictionary of data. """
    data = {}
    with open(file_path, 'r') as file:
        for line in file:
            frame_id, track_id, x, y, w, h, det_score, *rest = map(float, line.strip().split(','))
            frame_id, track_id = int(frame_id), int(track_id)
            if frame_id not in data:
                data[frame_id] = []
            data[frame_id].append((track_id, (x, y, w, h)))
    return data

def compute_iou(bbox1, bbox2):
    """ Compute the Intersection over Union (IoU) of two bounding boxes. """
    x1, y1, w1, h1 = bbox1
    x2, y2, w2, h2 = bbox2

    xi1 = max(x1, x2)
    yi1 = max(y1, y2)
    xi2 = min(x1 + w1, x2 + w2)
    yi2 = min(y1 + h1, y2 + h2)
    inter_area = max(xi2 - xi1, 0) * max(yi2 - yi1, 0)

    bbox1_area = w1 * h1
    bbox2_area = w2 * h2
    union_area = bbox1_area + bbox2_area - inter_area

    iou = inter_area / union_area if union_area != 0 else 0

    return iou

def calculate_mota_idf1(gt_file, pred_file):
    gt_data = read_tracking_file(gt_file)
    pred_data = read_tracking_file(pred_file)

    total_frames = max(max(gt_data.keys()), max(pred_data.keys()))
    total_misses, total_false_positives, total_mismatches, total_matches = 0, 0, 0, 0
    total_gt_detections, total_pred_detections = 0, 0

    for frame in range(total_frames + 1):
        gt_detections = gt_data.get(frame, [])
        pred_detections = pred_data.get(frame, [])

        # Create the cost matrix based on IoU
        cost_matrix = np.zeros((len(gt_detections), len(pred_detections)))
        for i, (_, gt_bbox) in enumerate(gt_detections):
            for j, (_, pred_bbox) in enumerate(pred_detections):
                cost_matrix[i, j] = -compute_iou(gt_bbox, pred_bbox)
        
        # Hungarian algorithm for optimal assignment
        gt_indices, pred_indices = linear_sum_assignment(cost_matrix)

        matched_ids = set()
        for g_idx, p_idx in zip(gt_indices, pred_indices):
            if cost_matrix[g_idx, p_idx] < -0.5:  # Threshold for IoU
                matched_ids.add(pred_detections[p_idx][0])
                total_matches += 1

        total_gt_detections += len(gt_detections)
        total_pred_detections += len(pred_detections)
        total_misses += len(gt_detections) - len(matched_ids)
        total_false_positives += len(pred_detections) - len(matched_ids)
        total_mismatches += sum([1 for _, pred_id in pred_detections if pred_id not in matched_ids])

    mota = 1 - (total_misses + total_false_positives + total_mismatches) / total_gt_detections
    idf1 = 2 * total_matches / (total_gt_detections + total_pred_detections)

    return mota, idf1

# File paths
prediction_file_path = './tracking_bytetrack/2023_12_28_17_39_28.txt'
gt_file_path = './tracking_gt/c5_gt.txt'

# Calculate MOTA and IDF1
mota, idf1 = calculate_mota_idf1(gt_file_path, prediction_file_path)
print(f"MOTA: {mota}, IDF1: {idf1}")