import time
import cv2
import numpy as np
import torch
from models.experimental import attempt_load
from utils.datasets import letterbox
from utils.general import check_img_size, non_max_suppression, scale_coords, xyxy2xywh, set_logging, check_requirements
from utils.plots import colors, Annotator  # plot_one_box
from utils.torch_utils import select_device  # time_synchronized
from grabscreen import grab_screen
from PIL import Image
import pyautogui
import math


pyautogui.FAILSAFE = False

pi = 3.1415916



@torch.no_grad()
def detect(
        # --------------------这里更改配置--------------------
        # ---------------------------------------------------
        weights='runs/train/exp40/weights/best.pt',  # 训练好的模型路径
        imgsz=(640,640),  # 训练模型设置的尺寸
        cap=0,  # 摄像头
        conf_thres=0.25,  # 置信度
        iou_thres=0.45,  # NMS IOU 阈值
        max_det=1000,  # 最大侦测的目标数
        device='0',  # 设备
        crop=True,  # 显示预测框
        classes=None,  # 种类
        agnostic_nms=False,  # class-agnostic NMS
        augment=False,  # 是否扩充推理
        half=False,  # 使用FP16半精度推理
        hide_labels=False,  # 是否隐藏标签
        hide_conf=False,  # 是否隐藏置信度
        line_thickness=3  # 预测框的线宽
):
    # #--------------------这里更改配置--------------------
    # -----------------------------------------------------

    # -----初始化-----
    set_logging()
    # 设置设备
    device = select_device(device)
    # CUDA仅支持半精度
    half &= device.type != 'cpu'

    # -----加载模型-----
    # 加载FP32模型
    model = attempt_load(weights, map_location=device)
    # 模型步幅
    stride = int(model.stride.max())
    # 检查图像大小
    imgsz = check_img_size(imgsz, s=stride)
    # 获取类名
    names = model.module.names if hasattr(model, 'module') else model.names
    # toFP16
    if half:
        model.half()
    mark = 0
    run_state = 1
    # -----进入循环：ESC退出-----
    picnum = 0
    while (True):
        image_array = grab_screen(region=(88, 92, 1914, 914))

        array_to_image = Image.fromarray(image_array, mode='RGB')  # 将array转成图像，才能送入yolo进行预测
        img = np.asarray(array_to_image)  # 将图像转成array
        # 设置labels--记录标签/概率/位置
        labels = []
        # 计时
        t0 = time.time()
        img0 = img
        # 填充调整大小
        img = letterbox(img0, imgsz, stride=stride)[0]


        # Convert
        img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
        img = np.ascontiguousarray(img)

        # # 转换
        # img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
        # img = np.ascontiguousarray(img)

        img = torch.from_numpy(img).to(device)
        # uint8 to fp16/32
        img = img.half() if half else img.float()
        # 0 - 255 to 0.0 - 1.0
        img /= 255.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        # 推断
        # t1 = time_synchronized()
        pred = model(img, augment=augment)[0]

        # 添加 NMS
        pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
        # t2 = time_synchronized()

        # 目标进程
        for i, det in enumerate(pred):  # 每幅图像的检测率
            s, im0 = '', img0.copy()
            # 输出字符串
            s += '%gx%g ' % img.shape[2:]
            # 归一化增益
            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]
            annotator = Annotator(im0, line_width=line_thickness, example=str(names))
            if len(det):
                # 将框从img_大小重新缩放为im0大小
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
                # 输出结果
                for c in det[:, -1].unique():
                    # 每类检测数
                    n = (det[:, -1] == c).sum()
                    # 添加到字符串
                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "
                    # 结果输出
                for *xyxy, conf, cls in reversed(det):
                    # 归一化xywh
                    xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()
                    # 标签格式
                    line = (cls, *xywh, conf)
                    # 整数类
                    c = int(cls)
                    # 建立标签
                    label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')

                    # 画预测框
                    if crop:
                        # print('right')
                        annotator.box_label(xyxy, label, color=colors(c, True))
                        # plot_one_box(xyxy, im0, label=label, color=colors(c, True), line_thickness=line_thickness)
                    # 记录标签/概率/位置
                    labels.append([names[c], conf, xyxy])

                    # 获取中心点，即分别求横、纵坐标的中间点
                    pointx = int((xyxy[0] + xyxy[2]) / 2)
                    pointy = int((xyxy[1] + xyxy[3]) / 2)

                    # # # 移动鼠标到中心点位置，并点击
                    # if run_state == 1:
                    if mark==0:
                        pyautogui.moveTo(1437, 662, _pause=False)
                        pyautogui.mouseDown()  # 鼠标按下
                        my_ts1 = time.time()
                        mark = 1
                    else:
                        #pyautogui.moveTo(1437, 662, _pause=False)
                        my_ts2 = time.time()
                        d2x = pointx-1826/2
                        d2y = pointy-822/2
                        a = math.atan2(d2y,d2x)
                        # d25x = (d2x*math.sin(a)+d2y*math.cos(a))/2*math.sin(a)*math.cos(a)
                        # d25y = (d2y*math.cos(a)-d2x*math.sin(a))/2*math.sin(a)*math.cos(a)
                        b=0
                        if 0<a<pi/2:
                            b = a- (pi/12)*(1-math.tan(abs(pi/4-a)))
                        if pi/2<=a<pi:
                            b = a- (pi/12)*(1-math.tan(abs(pi*3/4-a)))
                        if -pi<a<-pi/2:
                            b = a+ (pi/12)*(1-math.tan(abs(pi*3/4+a)))
                        if -pi/2<=a<=0:
                            b = a- (pi/12)*(1-math.tan(abs(pi/4+a)))

                        pyautogui.moveTo(1437+140*math.cos(b),\
                                         662+140*math.sin(b), _pause=False)
                        if (my_ts2-my_ts1)>3:
                            pyautogui.mouseUp()  # 鼠标释放
                            mark =0

            # # 显示图片
            imshow = cv2.cvtColor(im0, cv2.COLOR_BGR2BGRA)
            cv2.imshow("copy window", imshow)
            key = cv2.waitKey(2)
            # # 这里设置ESC退出
            # if key == 'q':
            #     run_state = ~run_state
            # 输出计算时间
            # print(f'消耗时间: ({time.time() - t0:.3f}s)')
        # --------------------END--------------------
        # -------------------------------------------------
    cv2.destroyAllWindows()


if __name__ == "__main__":
    detect()