import torch
import cv2
import os
import numpy as np
from pathlib import Path
from sort import Sort
from scipy.optimize import linear_sum_assignment

from models.common import DetectMultiBackend
from utils.augmentations import letterbox
from utils.general import non_max_suppression, scale_boxes
from models.experimental import attempt_load
from utils.torch_utils import select_device, smart_inference_mode

from tqdm import tqdm


def bbox_iou(box1, box2):
    x1_max = max(box1[0], box2[0])
    y1_max = max(box1[1], box2[1])
    x2_min = min(box1[2], box2[2])
    y2_min = min(box1[3], box2[3])

    intersection = max(0, x2_min - x1_max + 1) * max(0, y2_min - y1_max + 1)
    area_box1 = (box1[2] - box1[0] + 1) * (box1[3] - box1[1] + 1)
    area_box2 = (box2[2] - box2[0] + 1) * (box2[3] - box2[1] + 1)

    union = area_box1 + area_box2 - intersection
    return intersection / union


# Model
weights = '/data/yt/credbugs/1028/yolov5/best.pt'
device = select_device(0)
model = DetectMultiBackend(weights, device=device, dnn=False, data='data/coco128.yaml', fp16=False)
stride = model.stride  # model stride
names = model.module.names if hasattr(model, 'module') else model.names

# input video path
vid_path = 'test1.mp4'

# output video path
out_path = vid_path.split('.')[0]+'_result.mp4'

# capture video
cap = cv2.VideoCapture(vid_path)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(cap.get(cv2.CAP_PROP_FPS))

# output video writer
fourcc = cv2.VideoWriter_fourcc(*'MP4V')
out = cv2.VideoWriter(out_path, fourcc, fps, (width, height))

# Initialize SORT tracker
tracker = Sort()

# Create a text file 
if not os.path.exists('results'):
    os.makedirs('results')
track_results_file = open('results/track_results.txt', 'w')
det_results_file = open('results/det_results.txt', 'w')

count = 0

# get the total number of frames in the video
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
pbar = tqdm(total=total_frames, desc="Processing frames", unit="frame")

track_id_class_matching_dic = {}

while True:
    ret, frame = cap.read()
    if not ret:
        break

    # Resize and pad image while meeting stride-multiple constraints
    img = letterbox(frame, new_shape=640, stride=stride)[0]

    # Convert
    img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3xHxW
    img = np.ascontiguousarray(img)

    # Convert to torch tensor
    img = torch.from_numpy(img).to(device)
    img = img.float()  # uint8 to fp32
    img /= 255.0  # 0 - 255 to 0.0 - 1.0

    if img.ndimension() == 3:
        img = img.unsqueeze(0)

    # Inference
    pred = model(img, augment=False)[0]

    # Apply NMS
    pred = non_max_suppression(pred, 0.25, 0.45, classes=None, agnostic=False)

    # Initialize a list to hold the detections for SORT
    dets = []
    class_0_objects = []
    # Process detections
    for i, det in enumerate(pred):  # detections per image
        if len(det):
            # Rescale boxes from img_size to original size
            det[:, :4] = scale_boxes(img.shape[2:], det[:, :4], frame.shape).round()

            # Print results
            for *xyxy, conf, cls in reversed(det):        
                # class_name = names[int(cls)]
                x1 = int(xyxy[0].cpu().numpy())
                y1 = int(xyxy[1].cpu().numpy())
                x2 = int(xyxy[2].cpu().numpy())
                y2 = int(xyxy[3].cpu().numpy())
                score = round(float(f'{conf:.2f}'),2)
                if int(cls) == 0:
                    class_0_objects.append([x1, y1, x2, y2, score, int(cls)])
                else:
                    dets.append([x1, y1, x2, y2, score, int(cls)])

    # Sort class 0 objects based on score (in descending order)
    class_0_objects.sort(key=lambda x: x[4], reverse=True)
    
    # Retain class 0 for the highest scoring object and assign class 1 to others
    for i, obj in enumerate(class_0_objects):
        if i == 0:  # highest scoring object
            dets.append(obj)
        else:  # other class 0 objects
            obj[5] = 1  # change class to 1
            dets.append(obj)
                

    if len(dets) == 0:
        dets = np.empty((0, 5))
    else:
        dets = np.array(dets)
    # Update tracker
    trackers = tracker.update(dets)

    # Print results
    # for d in trackers:
    #     # print('%d,%d,%.2f,%.2f,%.2f,%.2f'%(count,d[4],d[0],d[1],d[2],d[3]))
    #     track_results_file.write('%d,%d,%.2f,%.2f,%.2f,%.2f\n' % (count, d[4], d[0], d[1], d[2], d[3]))
        # frame_n,id,x,y,x2,y2

    # # dets (x1,y1,x2,y2,score,class_id)
    # for d in dets:
    #     det_results_file.write('%d,%d,%.2f,%.2f,%.2f,%.2f\n' % (count, d[5], d[0], d[1], d[2], d[3]))
    new_bit = 0
    for d in trackers:
        if d[4] not in track_id_class_matching_dic.keys():
            new_bit+=1

    if new_bit>0:
        iou_matrix = np.zeros((len(trackers), len(dets)))
        for i, trac_obj in enumerate(trackers):
            for j, det_obj in enumerate(dets):
                iou_matrix[i, j] = bbox_iou(trac_obj[:4], det_obj[:4])
    
        cost_matrix = 1 - iou_matrix
        rows, cols = linear_sum_assignment(cost_matrix)
    
        trackers_list = trackers.tolist()
        for row, col in zip(rows, cols):
            if iou_matrix[row, col] > 0.3 and int(trackers_list[row][4]) not in track_id_class_matching_dic.keys():
                track_id_class_matching_dic[int(trackers_list[row][4])] = int(dets[col][5])
    
    # print(track_id_class_matching_dic)
    
    # Write frame to output video
    for trac_obj in trackers:
        track_id = int(trac_obj[4])
        class_id = track_id_class_matching_dic.get(track_id, None)
        if class_id is not None:  # Only draw if class_id is available
            x1, y1, x2, y2 = map(int, trac_obj[:4])
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)  # Green bounding box with thickness 2
            label = f'Track ID: {track_id}, Class ID: {class_id}'
            cv2.putText(frame, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)  # Blue text

    # Write the frame with bounding boxes to the output video
    out.write(frame)
    
    count+=1
    pbar.update(1)

# release
cap.release()
pbar.close()

track_results_file.close()
det_results_file.close()