from scipy.optimize import linear_sum_assignment

def iou(bb_test, bb_gt):
    """
    Computes IUO between two bboxes in the form [x1,y1,x2,y2]
    """
    xx1 = np.maximum(bb_test[0], bb_gt[0])
    yy1 = np.maximum(bb_test[1], bb_gt[1])
    xx2 = np.minimum(bb_test[2], bb_gt[2])
    yy2 = np.minimum(bb_test[3], bb_gt[3])
    w = np.maximum(0., xx2 - xx1)
    h = np.maximum(0., yy2 - yy1)
    wh = w * h
    o = wh / ((bb_test[2]-bb_test[0])*(bb_test[3]-bb_test[1])
              + (bb_gt[2]-bb_gt[0])*(bb_gt[3]-bb_gt[1]) - wh)
    return o

# create a dictionary to store the class_id for each tracking_id
tracking_class_ids = {}

# open the tracking results file
with open('results/track_results.txt', 'r') as track_file:
    # read the tracking results line by line
    for line in track_file:
        # split the line into a list
        track_data = line.strip().split(',')
        
        # parse the data
        frame_number, tracking_id, x1, y1, x2, y2 = map(float, track_data)
        frame_number = int(frame_number)
        tracking_id = int(tracking_id)
        
        # if this is the first frame for the tracking_id, find the class_id
        if tracking_id not in tracking_class_ids:
            track_bbox = [x1, y1, x2, y2]
            
            # open the detection results file
            with open('results/det_results.txt', 'r') as det_file:
                # read the detection results line by line
                det_bboxes = []
                det_classes = []
                for det_line in det_file:
                    # split the line into a list
                    det_data = det_line.strip().split(',')
                    
                    # parse the data
                    det_frame_number, class_id, det_x1, det_y1, det_x2, det_y2 = map(float, det_data)
                    det_frame_number = int(det_frame_number)
                    class_id = int(class_id)
                    
                    # if the frame number matches, calculate IoU
                    if det_frame_number == frame_number:
                        det_bbox = [det_x1, det_y1, det_x2, det_y2]
                        det_bboxes.append(det_bbox)
                        det_classes.append(class_id)
                
                # calculate IoU for each pair of track_bbox and det_bbox
                iou_matrix = np.zeros((1, len(det_bboxes)))
                for i, det_bbox in enumerate(det_bboxes):
                    iou_matrix[0, i] = iou(track_bbox, det_bbox)
                
                # apply Hungarian algorithm to find the best match
                row_ind, col_ind = linear_sum_assignment(-iou_matrix)
                if len(row_ind) > 0 and iou_matrix[row_ind[0], col_ind[0]] >= 0.5:
                    class_id = det_classes[col_ind[0]]
                    tracking_class_ids[tracking_id] = class_id

# now you can use the tracking_class_ids dictionary to find the class_id for each tracking_id
