import os
import shutil
import cv2
import numpy as np
from PyQt5.QtGui import QPixmap, QImage, QMouseEvent, QGuiApplication
from PyQt5.QtWidgets import QMessageBox, QFileDialog
from PyQt5.QtCore import QTimer, Qt
from PyQt5 import QtCore, QtGui , uic
import globf
from YoloClass import YoloThread

GLOBAL_STATE = True

formType, baseType = uic.loadUiType("./ui/win.ui")


class Mainwindow(formType, baseType):
    def __init__(self):
        super().__init__()
        # 加载UI
        self.setupUi(self)
        self.setWindowFlags(Qt.CustomizeWindowHint)

        # ui部件功能设置
        self.inputPath = ""
        # Slider
        self.con_slider.valueChanged.connect(self.ValueChange)
        self.iou_slider.valueChanged.connect(self.ValueChange)
        self.numcon = self.con_slider.value() / 100.0
        self.numiou = self.iou_slider.value() / 100.0

        # TOOLS
        self.folder.clicked.connect(self.Selectfile)
        self.importbtn.clicked.connect(self.Import)
        self.exporter.clicked.connect(self.Export)

        # 最大化 最小化 关闭
        # 最大化按钮图片变化
        MaxIcon = QtGui.QIcon()
        MaxIcon.addPixmap(QtGui.QPixmap("img/icons/square.png"), QtGui.QIcon.Normal, QtGui.QIcon.Off)
        MaxIcon.addPixmap(QtGui.QPixmap("img/icons/reduce.png"), QtGui.QIcon.Active, QtGui.QIcon.On)
        MaxIcon.addPixmap(QtGui.QPixmap("img/icons/reduce.png"), QtGui.QIcon.Selected, QtGui.QIcon.On)
        self.MaxButton.setCheckable(True)
        self.MaxButton.setIcon(MaxIcon)
        self.MaxButton.clicked.connect(self.max_or_restore)
        self.MinButton.clicked.connect(self.showMinimized)
        self.CloseButton.clicked.connect(self.close)
        # 视频预览
        self.input.setAlignment(QtCore.Qt.AlignCenter | QtCore.Qt.AlignVCenter)
        self.output.setAlignment(QtCore.Qt.AlignCenter | QtCore.Qt.AlignVCenter)

        # 视频操作
        # 播放按钮图片变化
        PlayIcon = QtGui.QIcon()
        PlayIcon.addPixmap(QtGui.QPixmap("img/icons/play-button.png"), QtGui.QIcon.Normal, QtGui.QIcon.Off)
        PlayIcon.addPixmap(QtGui.QPixmap("img/icons/pause.png"), QtGui.QIcon.Normal, QtGui.QIcon.On)
        PlayIcon.addPixmap(QtGui.QPixmap("img/icons/play-button.png"), QtGui.QIcon.Disabled, QtGui.QIcon.Off)
        PlayIcon.addPixmap(QtGui.QPixmap("img/icons/pause.png"), QtGui.QIcon.Disabled, QtGui.QIcon.On)
        PlayIcon.addPixmap(QtGui.QPixmap("img/icons/play-button.png"), QtGui.QIcon.Active, QtGui.QIcon.Off)
        PlayIcon.addPixmap(QtGui.QPixmap("img/icons/pause.png"), QtGui.QIcon.Active, QtGui.QIcon.On)
        PlayIcon.addPixmap(QtGui.QPixmap("img/icons/play-button.png"), QtGui.QIcon.Selected, QtGui.QIcon.Off)
        PlayIcon.addPixmap(QtGui.QPixmap("img/icons/pause.png"), QtGui.QIcon.Selected, QtGui.QIcon.On)
        self.playbtn.setCheckable(True)
        self.playbtn.setIcon(PlayIcon)

        # 自动加载 pt文件
        self.comboBox.clear()
        self.pt_Path = "./weights/best.pt"
        self.pt_list = os.listdir('./weights/')
        self.pt_list = [file for file in self.pt_list if file.endswith('.pt')]
        self.pt_list.sort(key=lambda x: os.path.getsize('./weights/' + x))
        self.comboBox.clear()
        self.comboBox.addItems(self.pt_list)
        self.qtimer_search = QTimer(self)
        self.qtimer_search.timeout.connect(lambda: self.search_pt())
        self.qtimer_search.start(2000)
        self.comboBox.currentTextChanged.connect(self.change_model)

        # yolov7 thread
        self.yolo_thread = YoloThread()
        # 获取模型
        self.model_type = self.comboBox.currentText()
        self.yolo_thread.weights = "./weights/%s" % self.model_type
        self.yolo_thread.percent_length = self.progressBar.maximum()
        self.yolo_thread.send_input.connect(lambda x: self.showimg(x, self.input, 'img'))
        self.yolo_thread.send_output.connect(lambda x: self.showimg(x, self.output, 'img'))
        self.yolo_thread.send_result.connect(self.show_result)
        self.yolo_thread.send_msg.connect(lambda x: self.foot_print(x))
        self.yolo_thread.send_percent.connect(lambda x: self.progressBar.setValue(x))
        # self.yolo_thread.send_fps.connect(lambda x: self.fps_label.setText(x))

        # 运行或停止 检测
        self.playbtn.clicked.connect(self.run_or_continue)
        self.stopbtn.clicked.connect(self.stop)


    # 寻找pt模型
    def search_pt(self):
        pt_list = os.listdir('./weights/')
        pt_list = [file for file in pt_list if file.endswith('.pt')]
        pt_list.sort(key=lambda x: os.path.getsize('./weights/' + x))

        if pt_list != self.pt_list:
            self.pt_list = pt_list
            self.comboBox.clear()
            self.comboBox.addItems(self.pt_list)

    # Conf 和 IoU 变化
    def ValueChange(self):
        self.numcon = self.con_slider.value() / 100.0
        self.numiou = self.iou_slider.value() / 100.0
        self.con_num.setValue(self.numcon)
        self.yolo_thread.conf = self.numcon
        self.iou_num.setValue(self.numiou)
        self.yolo_thread.iou = self.numiou

    # Model 变化
    def change_model(self, x):
        self.model_type = self.comboBox.currentText()
        self.yolo_thread.weights = "./weights/%s" % self.model_type

    # 显示Label图片
    @staticmethod
    def showimg(img, label, flag):
        try:
            if flag == "path":
                img_src = cv2.imdecode(np.fromfile(img, dtype=np.uint8), -1)
            else:
                img_src = img
            ih, iw, _ = img_src.shape
            w = label.geometry().width()
            h = label.geometry().height()
            # keep original aspect ratio
            if iw / w > ih / h:
                scal = w / iw
                nw = w
                nh = int(scal * ih)
                img_src_ = cv2.resize(img_src, (nw, nh))
            else:
                scal = h / ih
                nw = int(scal * iw)
                nh = h
                img_src_ = cv2.resize(img_src, (nw, nh))

            frame = cv2.cvtColor(img_src_, cv2.COLOR_BGR2RGB)
            img = QImage(frame.data, frame.shape[1], frame.shape[0], frame.shape[2] * frame.shape[1],
                         QImage.Format_RGB888)
            label.setPixmap(QPixmap.fromImage(img))

        except Exception as e:
            print(repr(e))

    # 选择照片/视频 并展示
    def Selectfile(self):
        file, _ = QFileDialog.getOpenFileName(
            self,  # 父窗口对象
            "Select image/video",  # 标题
            "./",  # 默认打开路径为当前路径
            "image/video type (*.jpg *.jpeg *.png *.bmp *.dib  *.jpe  *.jp2 *.mp4)"  # 选择类型过滤项，过滤内容在括号中
        )
        if file == "":
            pass
        else:
            self.inputPath = file
            globf.set_value('inputPath', self.inputPath)
            if ".avi" in self.inputPath or ".mp4" in self.inputPath:
                # 显示第一帧
                self.cap = cv2.VideoCapture(self.inputPath)
                ret, frame = self.cap.read()
                if ret:
                    rgbImage = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    self.showimg(rgbImage, self.input, 'img')
            else:
                self.showimg(self.inputPath, self.input, 'path')

    # 导入模块
    def Import(self):
        file, _ = QFileDialog.getOpenFileName(
            self,  # 父窗口对象
            "select import model",  # 标题
            "./",  # 默认打开路径为当前路径
            "model type (*.pt)"  # 选择类型过滤项，过滤内容在括号中
        )
        if file == "":
            pass
        else:
            shutil.copy(file, self.pt_Path)
            QMessageBox.information(self, 'Notice', 'import model successfully!')

    # 导出结果
    def Export(self):
        self.OutputDir, _ = QFileDialog.getSaveFileName(
            self,  # 父窗口对象
            "export image/video",  # 标题
            r".",  # 起始目录
            "image/video type (*.jpg *.jpeg *.png *.bmp *.dib  *.jpe  *.jp2 *.mp4)"  # 选择类型过滤项，过滤内容在括号中
        )
        if self.output == "":
            QMessageBox.warning(self, 'Notice', 'Choose save path')
        else:
            try:
                shutil.copy(self.yolo_thread.save_path, self.OutputDir)
                QMessageBox.warning(self, 'Notice', 'Export model successfully!!')
            except Exception as e:
                QMessageBox.warning(self, 'Notice', 'Finish the recognition first')
                print(e)

    # 最大化最小化窗口
    def max_or_restore(self):
        global GLOBAL_STATE
        status = GLOBAL_STATE
        if status:
            self.showMaximized()
            GLOBAL_STATE = False
        else:
            self.showNormal()
            GLOBAL_STATE = True

    # 开始/暂停 预测
    def run_or_continue(self):
        if self.inputPath == "":
            QMessageBox.warning(self, "Notice", "Select image/video first")
        else:
            self.yolo_thread.jump_out = False
            if self.playbtn.isChecked():
                self.yolo_thread.is_continue = True
                if not self.yolo_thread.isRunning():
                    self.yolo_thread.start()
                    self.foot_print("Begin identity >>>>>" + self.inputPath )
            else:
                self.yolo_thread.is_continue = False
                self.foot_print('Pause')

    # 停止识别
    def stop(self):
        self.yolo_thread.jump_out = True
        self.foot_print('Stop')

    # 统计结果
    def show_result(self, statistic_dic):
        try:
            self.resultlist.clear()
            statistic_dic = sorted(statistic_dic.items(), key=lambda x: x[1], reverse=True)
            statistic_dic = [i for i in statistic_dic if i[1] > 0]
            results = [' ' + str(i[0]) + '：' + str(i[1]) for i in statistic_dic]
            self.resultlist.addItems(results)

        except Exception as e:
            print(repr(e))

    # foot栏 输出结果
    def foot_print(self, msg):
        if msg in ['Stop','Finished']:
            self.playbtn.setChecked(False)
        self.outputbox.setText(msg)


class MyWindow(Mainwindow):
    def __init__(self):
        super(MyWindow, self).__init__()
        self.center()

    def mousePressEvent(self, event: QMouseEvent) -> None:
        if event.button() == Qt.LeftButton:
            self.mouse_start_pt = event.globalPos()
            self.window_pos = self.frameGeometry().topLeft()
            self.drag = True

    def mouseMoveEvent(self, event: QMouseEvent) -> None:
        if Qt.LeftButton and self.drag:
            distance = event.globalPos() - self.mouse_start_pt
            self.move(self.window_pos + distance)

    def mouseReleaseEvent(self, event: QMouseEvent) -> None:
        if event.button() == Qt.LeftButton:
            self.drag = False

    def center(self):
        # PyQt6获取屏幕参数
        screen = QGuiApplication.primaryScreen().size()
        size = self.geometry()
        self.move((screen.width() - size.width()) / 2,
                  (screen.height() - size.height()) / 2 - 10)
