import torch
import cv2
import numpy as np
from pathlib import Path
from PIL import Image, ImageDraw

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

####################################################
#################### Parameters ####################
# Define your object class indices and corresponding colors
## 사람 (Person): 0

# input video path
vid_path = 'nighttime.mp4'

# object detection threshold
conf_thres=0.25  # confidence threshold
iou_thres=0.45  # NMS IOU threshold
####################################################
####################################################


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

# Model
weights = './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

# 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, conf_thres, iou_thres, classes=None, agnostic=False)

    # Reset people count for each frame
    people_count = 0
    
    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 == 0:  # Check if the class is 0 (people)
                    people_count += 1  # Increment people count
                    # Draw red bounding box for each detected person
                    cv2.rectangle(frame, (int(xyxy[0]), int(xyxy[1])), (int(xyxy[2]), int(xyxy[3])), (0, 0, 255), 2)
        
    # Display people count on the frame
    cv2.putText(frame, f"People Counting: {people_count}", (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA)

    # Write frame to output video
    out.write(frame)
    pbar.update(1)

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