import datetime
import os
import time
import ffmpeg
import torch
import cv2
import numpy as np
from multiprocessing import Process, Manager
from threading import Thread
from read_data import LoadImages, LoadStreams
import torch.backends.cudnn as cudnn

class ElectromobileDetection():
    time_reference = datetime.datetime.now()
    counter_frame = 0
    processed_fps = 0

    def __init__(self,video_path=None):
        
        self.model = torch.hub.load((os.getcwd()) + "/algorithm/yolov5", 'custom', source='local', path='./weight/electromobile.pt', force_reload=True)
        self.classes = ["电动车", "摩托车"]
     
        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.dataset = LoadImages(self.video_name)
        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')
        source = source
        self.imgsz = 640
        cudnn.benchmark = True
        self.dataset = LoadStreams(source, img_size=self.imgsz)
        self.flag = 1

    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)
            # print(results)
            img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

        # Loop through each detected object and count the people
            num_people = 0
            bgr = (0, 255, 0)

            txt = ""
            objs = results.xyxy[0]
            for c in objs[:,-1].unique():
                n = (objs[:,-1] == c).sum()  # detections per class
                txt += f"{n} {self.classes[int(c)]}{'s' * (n > 1)}, "  # add to string
                
                
            for obj in objs:
                if obj[-1] == 0:  # 1 is the class ID for '未戴头盔'
                    # Draw bounding boxes around people
                    xmin, ymin, xmax, ymax = map(int, obj[:4])
                    accuracy = obj[4]
                    if (accuracy > 0.2):
                        cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 0, 255), 2)
                        cv2.putText(img, f" {round(float(accuracy), 2)}", (xmin, ymin),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 255), 2)
                
                elif obj[-1] == 1:  # 1 is the class ID for '未戴头盔'
                    # Draw bounding boxes around people
                    xmin, ymin, xmax, ymax = map(int, obj[:4])
                    accuracy = obj[4]
                    if (accuracy > 0.2):
                        cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)
                        cv2.putText(img, f" {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(jpeg.shape)
            
            return jpeg.tobytes(),  txt