import torch
import cv2
import numpy as np
from pathlib import Path

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

# Define your class indices and corresponding colors
obj_class = [0, 1, 2, 3, 5, 7, 25]
obj_color = [
    (0, 0, 255),
    (0, 255, 0),
    (255, 0, 0),
    (0, 255, 255),
    (255, 0, 255),
    (255, 255, 0),
    (128, 0, 0)
]

# Dictionary to map class indices to their colors
class_to_color = {cls_idx: color for cls_idx, color in zip(obj_class, obj_color)}

# Model
weights = '/data/yt/credbugs/1104/yolov5x.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 = 'daytime.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))


# 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")

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 = []
    # 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()
            
           # Process detections
            for *xyxy, conf, cls in det:
                cls_int = int(cls)
                if cls_int in obj_class:
                    # Get bounding box coordinates
                    x1 = int(xyxy[0].cpu().numpy())
                    y1 = int(xyxy[1].cpu().numpy())
                    x2 = int(xyxy[2].cpu().numpy())
                    y2 = int(xyxy[3].cpu().numpy())
    
                    # Get the color for the class
                    color = class_to_color[cls_int]
    
                    # Draw the bounding box
                    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
                    
    # Write frame to output video
    out.write(frame)
    pbar.update(1)

# release video capture and writer
cap.release()
out.release()
pbar.close()