import os

from PyQt5.QtCore import Qt
from PyQt5.QtGui import QPixmap, QMouseEvent, QKeySequence, QTransform, QImage
from mainFrame import Ui_MainWindow
import sys
from PyQt5.QtWidgets import QApplication, QMainWindow, QFileDialog, QMessageBox, QAction, QMenu, QInputDialog, \
    QShortcut, QTableWidgetItem, QGraphicsScene
from PyQt5 import QtGui, QtWidgets, QtCore
import cv2
import numpy as np

onnx = None
if os.path.exists('models/circle_model.onnx'):
    import onnxruntime as ort


    class Onnx:
        def __init__(self, root):
            self.ort_session = ort.InferenceSession(root)

        def to_rgb(self, im):
            """
            图片转 RGB
            :param im: 图片 np.array
            :return: 图片 np.array
            """
            if len(im.shape) == 2:  # 灰度
                return cv2.cvtColor(im, cv2.COLOR_GRAY2RGB)
            elif im.shape[2] == 4:  # PNG (RGBA)
                return cv2.cvtColor(im, cv2.COLOR_BGRA2RGB)
            assert len(im.shape) == 3 and im.shape[2] == 3  # 如果 Channel 不为3, 肯定出问题了
            return cv2.cvtColor(im, cv2.COLOR_BGR2RGB)

        def normalize(self, img):
            """
            使用 OpenCV 对图像进行归一化
            :param img: 图片 np.array
            :return: 归一化后的图片 np.array
            """
            mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
            std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
            img = img.astype(np.float32) / 255.0  # 确保图像是 float32 类型
            img = (img - mean) / std
            return img

        def __call__(self, img):
            img = cv2.resize(img, (128, 128))
            img = self.to_rgb(img)
            mask = np.zeros((128, 128), dtype=np.uint8)
            cv2.circle(mask, (64, 64), 62, 255, -1)
            img[mask != 255] = 0
            img = self.normalize(img)
            img = img.transpose(2, 0, 1)  # 将 Channel 放在前面
            img = np.expand_dims(img, axis=0)
            loc = self.ort_session.run(None, {'input': img})[0]
            loc = np.array(loc)
            angle = int(loc.argmax(1)[0])
            print('angle:', angle)
            return angle


    onnx = Onnx('models/circle_model.onnx')


def img_resize(image):
    height, width = image.shape[0], image.shape[1]
    # 设置新的图片分辨率框架
    width_new = 128
    height_new = 128
    # 判断图片的长宽比率
    if width / height >= width_new / height_new:
        img_new = cv2.resize(image, (width_new, int(height * width_new / width)))
    else:
        img_new = cv2.resize(image, (int(width * height_new / height), height_new))
    return img_new


class MainFrameShow(Ui_MainWindow):
    def __init__(self, main_window):
        super(MainFrameShow, self).setupUi(main_window)
        self.MainWindow = main_window
        self.import_path = None
        self.save_path = None
        # 已经保存的文件列表
        self.save_file_list = []
        self.graphicsView_list = []
        self.rotation_degree = 2
        self.col_num = 6
        self.row_num = 30
        self.init_ui_event()

    def init_ui_event(self):
        """
        初始化UI事件
        :return:
        """
        # 设置菜单栏
        # self.init_menu()
        # 设置快捷键
        self.init_shortcut()
        # 设置按钮事件
        self.init_button_event()
        # 设置滚动条事件
        # self.init_scroll_event()
        # 设置鼠标事件
        self.init_mouse_event()
        # 设置表格事件
        self.init_tabel_event()
        # 设置graphicsView事件
        self.init_horizontal_slider_event()

    def init_horizontal_slider_event(self):
        """
        设置horizontal_slider事件
        :return:
        """
        self.horizontalSlider.valueChanged.connect(self.horizontal_slider_value_changed)
        # 鼠标抬起事件
        self.horizontalSlider.sliderReleased.connect(self.horizontal_slider_released)

    def horizontal_slider_value_changed(self):
        """
        horizontal_slider值改变事件,主要是展示图片
        :return:
        """
        # 计算点击图片的度数
        degree = int(self.horizontalSlider.value())
        # 获取选中的行
        selected_items = self.files_table.selectedItems()
        if len(selected_items) == 0:
            return
        # 取文件名
        row = selected_items[0].row()
        file_name = self.files_table.item(row, 0).text()
        # 把图片旋转
        pic_path = os.path.join(self.import_path, file_name)
        self.show_pic_to_rotate(pic_path, degree)
        self.degree_label.setText(f"{degree}")

    def horizontal_slider_released(self):
        """
        horizontal_slider鼠标抬起事件
        :return:
        """
        self.save_pic_with_degree()

    def save_pic_with_degree(self):
        """
        带度数保存图片
        :return:
        """
        # 计算点击图片的度数
        degree = int(self.horizontalSlider.value())
        # 获取选中的行
        selected_items = self.files_table.selectedItems()
        if len(selected_items) == 0:
            return
        # 取文件名
        row = selected_items[0].row()
        file_name = self.files_table.item(row, 0).text()
        # 保存的文件名为原文件名+度数.原后缀
        save_file_name = f'{file_name.split(".")[0]}|{degree}.{file_name.split(".")[-1]}'
        # 删除这张图片已经处理好的图片，命名规则为原文件名+度数.原后缀，度数不确定
        for save_file in self.save_file_list:
            if file_name.split('.')[0] in save_file:
                os.remove(os.path.join(self.save_path, save_file))
                self.save_file_list.remove(save_file)
                break
        # 把图片旋转
        pic_path = os.path.join(self.import_path, file_name)
        # pic = QPixmap(pic_path)
        # pixmap = pic.transformed(QTransform().rotate(degree))
        # # 保存图片
        save_pic_path = os.path.join(self.save_path, save_file_name)
        # # 保存图片
        # pixmap.save(save_pic_path)
        img = cv2.imread(pic_path)
        # 顺时针旋转degree度
        M = cv2.getRotationMatrix2D((img.shape[1] / 2, img.shape[0] / 2), -degree, 1)
        img = cv2.warpAffine(img, M, (img.shape[1], img.shape[0]))
        cv2.imwrite(save_pic_path, img)
        # 更新表格
        self.files_table.setItem(row, 1, QTableWidgetItem("已处理"))
        self.files_table.item(row, 1).setBackground(QtGui.QColor(0, 255, 0))
        # 更新已经保存的文件列表
        self.save_file_list.append(save_file_name)

    def init_shortcut(self):
        """
        设置快捷键
        :return:
        """
        # 设置热键s
        # self.hot_key_s = QAction()
        sq = QShortcut(QKeySequence('s'), self.MainWindow)  # 通过QShortcut控件设置快捷方式，触发槽函数
        sq.activated.connect(self.next_img)
        # 设置热键w
        wq = QShortcut(QKeySequence('w'), self.MainWindow)  # 通过QShortcut控件设置快捷方式，触发槽函数
        wq.activated.connect(self.prev_img)

    def init_button_event(self):
        """
        设置按钮事件
        :return:
        """
        self.import_pics_btn.clicked.connect(self.import_pics)
        self.set_pics_save_path_btn.clicked.connect(self.set_pics_save_path)

    def init_mouse_event(self):
        """
        设置鼠标事件
        :return:
        """
        pass

    def init_tabel_event(self):
        """
        设置表格事件
        :return:
        """
        # 设置选中内容改变事件
        self.files_table.itemSelectionChanged.connect(self.table_item_selected)

    def next_img(self):
        """
        下一张图片
        :return:
        """
        # 获取选中的行
        selected_items = self.files_table.selectedItems()
        if len(selected_items) == 0:
            return
        # 获取选中的行号
        row = selected_items[0].row()
        # 下一行
        row += 1
        if row >= self.files_table.rowCount():
            return
        self.files_table.selectRow(row)

    def prev_img(self):
        """
        上一张图片
        :return:
        """
        # 获取选中的行
        selected_items = self.files_table.selectedItems()
        if len(selected_items) == 0:
            return
        # 获取选中的行号
        row = selected_items[0].row()
        # 上一行
        row -= 1
        if row < 0:
            return
        self.files_table.selectRow(row)

    def table_item_selected(self):
        """
        表格选中内容改变事件
        :return:
        """
        # 获取选中的行
        selected_items = self.files_table.selectedItems()
        if len(selected_items) == 0:
            return
        # 度数轨道设置为 0
        self.horizontalSlider.setValue(0)
        # 度数标签设置为 0
        self.degree_label.setText(f"{0}")
        # 获取选中的行号
        row = selected_items[0].row()
        # 获取文件名
        file_name = self.files_table.item(row, 0).text()
        # 获取文件路径
        file_path = os.path.join(self.import_path, file_name)
        degree = 0
        # 检查保存文件中是否有
        for save_file in self.save_file_list:
            if file_name.split('.')[0] in save_file:
                # 获取度数
                degree = int(save_file.split('|')[1].split('.')[0])
                self.horizontalSlider.setValue(degree)
                break
        # 显示图片
        if degree == 0:
            if onnx is not None:
                # 会进行一次初次推理
                img = np.fromfile(file_path, np.uint8)
                img = cv2.imdecode(img, cv2.IMREAD_COLOR)
                degree = onnx(img)
                self.horizontalSlider.setValue(degree)
                # 保存图片
        self.save_pic_with_degree()
        self.show_pic_to_rotate(file_path, degree)

    def convert_cv2_to_qt(self, cv_img):
        """ Convert from an opencv image to QPixmap """
        # Convert the image from BGR to RGB format
        rgb_image = cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB)

        # Get the shape of the image
        height, width, channels = rgb_image.shape

        # Create QImage from the numpy array
        bytes_per_line = channels * width
        qt_image = QImage(rgb_image.data, width, height, bytes_per_line, QImage.Format_RGB888)

        # Convert QImage to QPixmap
        qt_pixmap = QPixmap.fromImage(qt_image)

        return qt_pixmap

    def show_pic_to_rotate(self, pic_path, degree=0):
        """
        显示图片旋转后的样子
        :param pic_path: 图片路径
        :param degree: 旋转度数
        :return:
        """
        self.degree_label.setText(f"{degree}")
        # 旋转图片
        img = cv2.imread(pic_path)
        # 顺时针旋转degree度
        M = cv2.getRotationMatrix2D((img.shape[1] / 2, img.shape[0] / 2), -degree, 1)
        img = cv2.warpAffine(img, M, (img.shape[1], img.shape[0]))
        pixmap = self.convert_cv2_to_qt(img)
        # pixmap = pic.transformed(QTransform().rotate(degree))
        # 设置图片
        scene = QGraphicsScene()
        # graphicsView.setScene(scene)
        pixmap_item = scene.addPixmap(pixmap)
        self.graphicsView.setScene(scene)
        self.graphicsView.fitInView(pixmap_item, Qt.KeepAspectRatio)
        QApplication.processEvents()
        # self.pic_label.setPixmap(pic)
        # self.pic_label.setAlignment(Qt.AlignCenter)
        # QApplication.processEvents()

    def import_pics(self):
        """
        导入图片
        :return:
        """
        if not self.save_path:
            QMessageBox.information(self.MainWindow, "提示", "请先设置图片保存路径")
            return
        dir_path = QFileDialog.getExistingDirectory(self.MainWindow, "选择图片读取目录", "./")

        if dir_path:
            self.import_path = dir_path
            # 加入文件表格
            self.init_files_table()

    def init_files_table(self):
        """
        初始化文件列表表格
        :return:
        """
        # 清空表格
        self.files_table.clearContents()
        self.files_table.setRowCount(0)
        files = os.listdir(self.import_path)
        files.sort()
        for file in files:
            if not (file.endswith(".jpg") or file.endswith(".png") or file.endswith(".jpeg") or file.endswith(".bmp")):
                continue
            # 添加到表格最后一行
            row = self.files_table.rowCount()
            self.files_table.insertRow(row)
            # 添加一行 文件名和未处理
            self.files_table.setItem(row, 0, QTableWidgetItem(file))
            # 判断是否已经处理，即file是否在save_file_list中，名字中包含即可
            self.files_table.setItem(row, 1, QTableWidgetItem("未处理"))
            for save_file in self.save_file_list:
                if file.split('.')[0] in save_file:
                    self.files_table.setItem(row, 1, QTableWidgetItem("已处理"))
                    # 设置当前行为绿色
                    self.files_table.item(row, 1).setBackground(QtGui.QColor(0, 255, 0))
                    break

    def set_pics_save_path(self):
        """
        设置图片保存路径
        :return:
        """
        save_path = QFileDialog.getExistingDirectory(self.MainWindow, "选择图片保存目录", "./")
        if save_path:
            self.save_path = save_path
            # 读取已经保存的文件列表
            self.read_save_file_list()

    def read_save_file_list(self):
        """
        读取已经保存的文件列表
        :return:
        """
        self.save_file_list = []
        save_file_list = os.listdir(self.save_path)
        for file in save_file_list:
            print(file)
            if not (file.endswith(".jpg") or file.endswith(".png") or file.endswith(".jpeg") or file.endswith(".bmp")):
                continue
            self.save_file_list.append(file)


if __name__ == '__main__':
    app = QApplication(sys.argv)
    MainWindow = QMainWindow()
    # 设置MainWindow不能调整大小
    MainWindow.setFixedSize(1024, 768)
    ui = MainFrameShow(MainWindow)
    MainWindow.show()
    sys.exit(app.exec_())