import os
import numpy as np
import motmetrics as mm
import copy

def read_results(filename, data_type, is_gt=False, is_ignore=False):
    """ Read tracking or ground truth data from a file. """
    results = {}
    with open(filename, 'r') as file:
        for line in file:
            line = line.strip()
            if line:
                frame_id, track_id, x, y, w, h, *rest = line.split(',')
                frame_id, track_id = int(frame_id), int(track_id)
                if frame_id not in results:
                    results[frame_id] = []
                if is_gt:
                    # For ground truth data
                    results[frame_id].append((track_id, (float(x), float(y), float(w), float(h))))
                elif not is_ignore:
                    # For tracking data
                    results[frame_id].append((track_id, (float(x), float(y), float(w), float(h))))
    return results

def unzip_objs(objs):
    """ Unzip object ids and bounding boxes. """
    ids, bboxes = [], []
    for obj_id, bbox in objs:
        ids.append(obj_id)
        bboxes.append(bbox)
    return ids, bboxes

class Evaluator(object):
    def __init__(self, data_root, seq_name, data_type):
        self.data_root = data_root
        self.seq_name = seq_name
        self.data_type = data_type

    def reset_accumulator(self):
        """ Reset the accumulator for a new evaluation. """
        self.acc = mm.MOTAccumulator(auto_id=True)

    def eval_frame(self, frame_id, trk_tlwhs, trk_ids, rtn_events=False):
        """ Evaluate a single frame. """
        # Ground truth and tracking data for the frame
        gt_objs = self.gt_frame_dict.get(frame_id, [])
        # gt_tlwhs, gt_ids = unzip_objs(gt_objs)[:2]
        gt_ids, gt_tlwhs = unzip_objs(gt_objs)[:2]

        # print(gt_tlwhs)
        # print(gt_ids)
        # print(trk_tlwhs)
        # print(trk_ids)
        
        # Calculate IoU distances and update the accumulator
        iou_distance = mm.distances.iou_matrix(gt_tlwhs, trk_tlwhs, max_iou=0.5)
        self.acc.update(gt_ids, trk_ids, iou_distance)

        if rtn_events and iou_distance.size > 0 and hasattr(self.acc, 'last_mot_events'):
            return self.acc.last_mot_events
        else:
            return None

    def eval_file(self, gt_filename, trk_filename):
        """ Evaluate an entire file. """
        self.reset_accumulator()
        self.gt_frame_dict = read_results(gt_filename, self.data_type, is_gt=True)
        result_frame_dict = read_results(trk_filename, self.data_type, is_gt=False)

        frames = sorted(list(set(self.gt_frame_dict.keys()) | set(result_frame_dict.keys())))
        for frame_id in frames:
            trk_objs = result_frame_dict.get(frame_id, [])
            trk_ids, trk_tlwhs = unzip_objs(trk_objs)[:2]
            # trk_tlwhs, trk_ids = unzip_objs(trk_objs)[:2]
            self.eval_frame(frame_id, trk_tlwhs, trk_ids, rtn_events=False)

        return self.acc

    @staticmethod
    def get_summary(accs, names, metrics=('mota', 'num_switches', 'idp', 'idr', 'idf1', 'precision', 'recall')):
        """ Get the summary of the evaluation. """
        if metrics is None:
            metrics = mm.metrics.motchallenge_metrics
        mh = mm.metrics.create()
        summary = mh.compute_many(
            accs, metrics=metrics, names=names, generate_overall=True
        )
        return summary

    @staticmethod
    def save_summary(summary, filename):
        """ Save the summary to an Excel file. """
        import pandas as pd
        writer = pd.ExcelWriter(filename)
        summary.to_excel(writer)
        writer.save()

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

# Example usage
evaluator = Evaluator('', 'seq_name', 'mot')
acc = evaluator.eval_file(gt_file_path, prediction_file_path)
summary = Evaluator.get_summary([acc], ['seq_name'])
print(summary)