import argparse
import os
import sys
from pathlib import Path

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

import mss
import numpy as np
import threading
from queue import Queue


FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

from models.common import DetectMultiBackend
from utils.datasets import IMG_FORMATS, VID_FORMATS, LoadImages, LoadStreams
from utils.general import (LOGGER, check_file, check_img_size, check_imshow, check_requirements, colorstr,
                           increment_path, non_max_suppression, print_args, scale_coords, strip_optimizer, xyxy2xywh)
from utils.plots import Annotator, colors, save_one_box
from utils.torch_utils import select_device, time_sync
from utils.augmentations import Albumentations, augment_hsv, copy_paste, letterbox, mixup, random_perspective

global detFifo, sizeList
sizeList = [int(1080/2 - 1080 / 8), int(1920 / 2 - 1920 / 8), int(1920/4), int(1080/4)]
detFifo = Queue(500)

def put_bbox(fifo, names, det):
    temp = []
    for *xyxy, conf, cls in reversed(det):
        c = int(cls)  # integer class
        label = names[c]
        dict = {'label': label, 'x0' : int(xyxy[0]), 'y0' : int(xyxy[1]), 'x1' : int(xyxy[2]), 'y1' : int(xyxy[3])}
        temp.append(dict)
    if temp != [] and (fifo.full() == False):
        fifo.put(temp)
class threading_detect(threading.Thread):
    def __init__(self, threadID, name):
        threading.Thread.__init__(self)
        self.threadID = threadID
        self.name = name
        self.monitor = {'top':sizeList[0], 'left':sizeList[1], 'width':sizeList[2], 'height':sizeList[3]}
    @torch.no_grad()
    def run(self, 
            #weights=ROOT / 'pretrained/yolov5s.pt',  # model.pt path(s)
            weights=ROOT / 'pretrained/yolov5Cs16.pt',
            imgsz=[640, 640],  # inference size (pixels)
            conf_thres=0.25,  # confidence threshold
            iou_thres=0.45,  # NMS IOU threshold
            max_det=1000,  # maximum detections per image
            device=0,  # cuda device, i.e. 0 or 0,1,2,3 or cpu
            view_img=True,  # show results
            classes=None,  # filter by class: --class 0, or --class 0 2 3
            agnostic_nms=False,  # class-agnostic NMS
            augment=False,  # augmented inference
            line_thickness=3,  # bounding box thickness (pixels)
            hide_labels=False,  # hide labels
            hide_conf=False,  # hide confidences
            half=False,  # use FP16 half-precision inference
            dnn=False,  # use OpenCV DNN for ONNX inference
            ):
        

        


        # Load model
        device = select_device(device)
        model = DetectMultiBackend(weights, device=device, dnn=dnn)
        stride, names, pt, jit, onnx = model.stride, model.names, model.pt, model.jit, model.onnx
        imgsz = check_img_size(imgsz, s=stride)  # check image size

        # Half
        half &= pt and device.type != 'cpu'  # half precision only supported by PyTorch on CUDA
        if pt:
            model.model.half() if half else model.model.float()

        cudnn.benchmark = True  # set True to speed up constant image size inference

        # Dataloader
                
        sct = mss.mss()
        if view_img:
            cv2.namedWindow('test')
        while(True):
            img0 = cv2.cvtColor(np.array(sct.grab(self.monitor)), cv2.COLOR_BGRA2BGR)  
            # Padded resize
            img = letterbox(img0, imgsz, stride=stride, auto=pt and not jit)[0]

            # Stack
            img = np.stack(img, 0)

            # Convert
            if len(img.shape) == 3:
                img = img[None]  # expand for batch dim    
            img = img[..., ::-1].transpose((0, 3, 1, 2))
            #img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = np.ascontiguousarray(img)

            bs = 1  # batch_size
            vid_path, vid_writer = [None] * bs, [None] * bs

            # Run inference
            if pt and device.type != 'cpu':
                model(torch.zeros(1, 3, *imgsz).to(device).type_as(next(model.model.parameters())))  # warmup
            dt, seen = [0.0, 0.0, 0.0], 0

            t1 = time_sync()
            im = torch.from_numpy(img).to(device)
            im = im.half() if half else im.float()  # uint8 to fp16/32
            im /= 255  # 0 - 255 to 0.0 - 1.0

            t2 = time_sync()
            dt[0] += t2 - t1

            # Inference
            pred = model(im, augment=augment, visualize=False)
            t3 = time_sync()
            dt[1] += t3 - t2

            # NMS
            pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
            dt[2] += time_sync() - t3

            # Second-stage classifier (optional)
            # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)

            # Process predictions
            for i, det in enumerate(pred):  # per image
                seen += 1
                im0 = img0

                
                gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
                imc = im0  
                annotator = Annotator(im0, line_width=line_thickness, example=str(names))
                if len(det):
                    # Rescale boxes from img_size to im0 size
                    det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()

                    # Print results
                    for c in det[:, -1].unique():
                        n = (det[:, -1] == c).sum()  # detections per class


                    # Write results
                    for *xyxy, conf, cls in reversed(det):
                        if view_img:  # Add bbox to image
                            c = int(cls)  # integer class
                            label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
                            annotator.box_label(xyxy, label, color=colors(c, True))
                            #LOGGER.info(f'label:{label} pos{[int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3])]}')
                    # Put Fifo
                    put_bbox(detFifo, names, det)
                # Print time (inference-only)
                temp = t3 - t2
                if temp <= 0.001:
                    temp = 0.001
                LOGGER.info(f'{np.int(1/(temp)) : d}FPS ({t3 - t2:.3f}s)')

                # Stream results
                im0 = annotator.result()
                if view_img:
                    cv2.imshow('test', im0)
                    cv2.waitKey(1)
            if cv2.waitKey(1) == ord('q'):  # q to quit
                break
        cv2.destroyAllWindows()  
def main():
    threadDetect = threading_detect(1, 'mssDetect')
    threadDetect.start()
    threadDetect.join()

if __name__ == "__main__":
    main()