import argparse
import os
import platform
import sys
from pathlib import Path

import cv2
import torch
import torch.backends.cudnn as cudnn

from read_data import LoadImages, LoadStreams

class TrafficDetection():
    def __init__(self, video_path=None, model=None):
        self.model = model
        self.classes = self.model.names
        self.imgsz = 640

        self.frame = [None]

        if video_path is not None:
            self.video_name = video_path
        else:
            self.video_name = 'vid2.mp4'  # A default video file

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.dataset = LoadImages(self.video_name, img_size=self.imgsz)
        self.flag = 0

    def use_webcam(self, source):
        # self.dataset.release()  # Release any existing video capture
        # self.cap = cv2.VideoCapture(0)  # Open default webcam
        # print('use_webcam')
        self.source = source
        cudnn.benchmark = True
        self.dataset = LoadStreams(source, img_size=self.imgsz)

    def class_to_label(self, x):
        return self.classes[int(x)]

    def get_frame(self):

        for im0s in self.dataset:
            # print(self.dataset.mode)
            # print(self.dataset)
            if self.dataset.mode == 'stream':
                img = im0s[0].copy()
            else:
                img = im0s.copy()

            img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
            results = self.model(img, size=640)
            img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

            # Loop through each detected object and count the people
            accuracy = 0
            num_people = 0

            color = (255, 200, 90)

            for obj in results.xyxy[0]:
                # xmin, ymin, xmax, ymax = map(int, obj[:4])
                # accuracy = obj[4]
                # if (accuracy > 0.5):

                #     cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 0, 255), 2)
                #     cv2.putText(img, f" {round(float(accuracy), 2), self.classes[obj[-1].item()]}", (xmin, ymin),
                #                     cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 255), 2)
                xmin, ymin, xmax, ymax = map(int, obj[:4])

                accuracy = obj[4]

                c = int(obj[-1])

                cv2.rectangle(img, (xmin, ymin), (xmax, ymax), color, 2)
                cv2.putText(img, f"{self.classes[c]}, {round(float(accuracy), 2)}", (xmin, ymin),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 255), 2)

        # Draw the number of people on the frame and display it
        ret, jpeg = cv2.imencode(".jpg", img)
        # print(num_people)
        return jpeg.tobytes(), num_people\