import os
import shutil
import sys

from imagededup.methods import DHash
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *


class WindowClass(QMainWindow):

    FILESPERPAGE = 20
    IMAGETYPE = ['jpg', 'jpeg', 'png', 'gif', 'bmp']
    TAGCOLOR = ["#FF002B", "#F78606", "#6200FF", "#FF571A",
                "#0095FF", "#0CD65D", "#FF0084", "#FFC400",
                "#7F7F7F", "#B312A2", "#6DC526", "#AC9802"]

    def __init__(self):
        super().__init__()
        # 设置和获取属性
        self.files_per_page = self.FILESPERPAGE
        self.image_path = None
        self.tag_list = []
        self.file_list = []
        self.file_count = len(self.file_list)
        self.page_num_total = (self.file_count + self.files_per_page - 1) \
            // self.files_per_page
        # 设置参数
        self.setWindowTitle('数据筛选工具')
        self.resize(1300, 800)
        # 初始化工具栏
        self._init_tool_bar()
        # 初始化主界面
        self._init_main_face()
        # 初始化停靠窗
        self._init_dock_window()
        # 快捷键
        self._init_hot_keys()

    #--------------------------------------#
    #                初始化部分              #
    #--------------------------------------#

    def _init_tool_bar(self):
        # 初始化工具栏
        tool = self.addToolBar("tools")
        tool.setStyleSheet("QToolBar{spacing:10px;}")
        # 打开按钮
        open = QAction(QApplication.style().standardIcon(
            QStyle.SP_FileDialogNewFolder), "open", self)
        open.triggered.connect(self.open_folder)
        tool.addAction(open)
        # 导出按钮
        export = QAction(QApplication.style().standardIcon(
            QStyle.SP_DialogHelpButton), "export", self)
        export.triggered.connect(self.export_window)
        tool.addAction(export)
        # 粗筛按钮
        filter = QAction(QApplication.style().standardIcon(
            QStyle.SP_DialogResetButton), "filter", self)
        filter.triggered.connect(self.remove_duplicate_window)
        tool.addAction(filter)
        # 上按钮
        pre = QAction(QApplication.style().standardIcon(
            QStyle.SP_ArrowBack), "pre", self)
        tool.addAction(pre)
        # 下按钮
        next = QAction(QApplication.style().standardIcon(
            QStyle.SP_ArrowForward), "next", self)
        tool.addAction(next)
        # 状态栏
        self.statusbar = self.statusBar()

    def _init_main_face(self):
        '''
        分页加载图片，减少内存占用
        '''
        # 分类面板
        # 选择类别
        self.tag_check_box_container = QListWidget()
        self.tag_check_box_container.setViewMode(QListView.IconMode)
        # 按钮
        self.tag_ok_button = QPushButton("确定")
        self.tag_new_button = QPushButton("添加")
        self.tag_new_button.clicked.connect(self.add_tag_item)
        # 汇总
        tag_face_layout = QHBoxLayout()
        tag_face_layout.addWidget(self.tag_check_box_container)
        tag_face_layout.addWidget(self.tag_ok_button)
        tag_face_layout.addWidget(self.tag_new_button)
        tag_face_container = QWidget()
        tag_face_container.setLayout(tag_face_layout)
        tag_face_container.setFixedHeight(50)

        # 图片显示
        self.file_list_widget = QListWidget()
        self.file_list_widget.setViewMode(QListView.IconMode)
        self.file_list_widget.setResizeMode(QListView.Adjust)
        # 高需要足够大，来满足缩放条件；如果item的高达到icon的上限，增大宽度的操作就会无效
        self.file_list_widget.setIconSize(QSize(300, 2000))
        self.file_list_widget.setDragDropMode(QAbstractItemView.NoDragDrop)
        self.file_list_widget.setSelectionMode(
            QAbstractItemView.ExtendedSelection)
        self.file_list_widget.clicked.connect(
            lambda i: self.show_single_image(self.index_shift + i.row())
            if 'single_image_win' in self.__dict__ and
            self.single_image_win.isVisible() else None)
        self.file_list_widget.doubleClicked.connect(
            lambda i: self.show_single_image(self.index_shift + i.row()))
        self.file_list_widget.contextMenuEvent = self.file_list_context_menu_event
        self.file_list_widget.wheelEvent = self.zoom_request
        self.file_list_widget.selectionChanged = self.get_files_tag

        # 页面切换
        self.first_page_button = QPushButton("首页")
        self.first_page_button.clicked.connect(lambda: self.fresh_page(1))
        self.pre_page_button = QPushButton("上一页")
        self.pre_page_button.clicked.connect(
            lambda: self.fresh_page(self.page_num - 1))
        self.page_to_line_edit = QLineEdit()  # 显示当前页数和执行页码跳转
        self.page_to_line_edit.setValidator(
            QRegExpValidator(QRegExp("[1-9][0-9]+$")))  # 仅允许数字输入
        self.page_to_line_edit.setFixedWidth(50)
        self.page_to_line_edit.returnPressed.connect(
            lambda: self.fresh_page(int(self.page_to_line_edit.text())))
        self.page_label = QLabel()  # 用于显示总页数
        self.page_label.setFixedWidth(50)
        self.next_page_button = QPushButton("下一页")
        self.next_page_button.clicked.connect(
            lambda: self.fresh_page(self.page_num + 1))
        self.last_page_button = QPushButton("尾页")
        self.last_page_button.clicked.connect(
            lambda: self.fresh_page(self.page_num_total))
        # 汇总
        self.page_control_layout = QHBoxLayout()
        self.page_control_layout.setContentsMargins(300, 0, 300, 0)
        self.page_control_layout.addWidget(self.first_page_button)
        self.page_control_layout.addWidget(self.pre_page_button)
        self.page_control_layout.addWidget(self.page_to_line_edit)
        self.page_control_layout.addWidget(self.page_label)
        self.page_control_layout.addWidget(self.next_page_button)
        self.page_control_layout.addWidget(self.last_page_button)
        self.page_control_container = QWidget()
        self.page_control_container.setLayout(self.page_control_layout)

        # 并入中心部件
        self.central_layout = QVBoxLayout()
        self.central_layout.setObjectName("central_layout") # NOTE 方便设置主题，本版本暂未加入主题功能
        self.central_layout.addWidget(tag_face_container)
        self.central_layout.addWidget(self.file_list_widget)
        self.central_layout.addWidget(self.page_control_container)
        self.central_widget = QWidget()
        self.central_widget.setObjectName("central_widget") # NOTE 方便设置主题，本版本暂未加入主题功能
        self.central_widget.setLayout(self.central_layout)

        self.setCentralWidget(self.central_widget)

    def _init_dock_window(self):
        # 文件目录系统窗口
        self.file_system_model = QFileSystemModel()
        self.file_system_model.setRootPath(QDir.currentPath())
        self.file_tree_view = QTreeView()
        self.file_tree_view.setModel(self.file_system_model)
        file_system_dock = QDockWidget("资源管理器")
        file_system_dock.setWidget(self.file_tree_view)

        # 统计窗口
        self.statics_tree = QTreeWidget()
        self.statics_tree.setColumnCount(2)
        self.statics_tree.setHeaderLabels(["类别", "数量"])
        self.statics_tree.setSortingEnabled(True)
        self.statics_tree.setSelectionMode(QAbstractItemView.ExtendedSelection)
        statics_dock = QDockWidget("批处理")
        statics_dock.setWidget(self.statics_tree)

        # 操作记录窗口
        self.operation_record = QTableWidget()
        operation_record_dock = QDockWidget("操作记录")
        operation_record_dock.setWidget(self.operation_record)
        operation_record_dock.setFeatures(QDockWidget.DockWidgetFloatable)

        # 加入主窗口
        # self.addDockWidget(Qt.DockWidgetArea.LeftDockWidgetArea, file_system_dock)
        self.addDockWidget(Qt.DockWidgetArea.RightDockWidgetArea, statics_dock)
        self.addDockWidget(
            Qt.DockWidgetArea.RightDockWidgetArea, operation_record_dock)

    def _init_hot_keys(self):
        '''
        设置快捷键
        '''
        # 翻页
        QShortcut("a", self, self.pre_page_button.click)
        QShortcut("d", self, self.next_page_button.click)
        # 类别设置快捷键
        for i in range(10):
            QShortcut(str(i), self, lambda i=i: self.set_tag_checkbox((i+9) % 10))
        # 全选快捷键
        QShortcut("Ctrl+a", self, self.file_list_widget.selectAll)

    #--------------------------------------#
    #               主窗口事件               #
    #--------------------------------------#
    # TODO 可以移除，每次打标签时自动保存
    def closeEvent(self, a0):
        '''
        退出时自动保存
        '''
        self.save_files_tag()
        super().closeEvent(a0)

    #--------------------------------------#
    #          标签选框窗口基础逻辑           #
    #--------------------------------------#
    def fresh_tag_items(self):
        '''
        根据tag_list更新
        '''
        self.tag_check_box_container.clear()
        for i, tag in enumerate(self.tag_list):
            item = QListWidgetItem(" "*(len(tag)+10))  # 使用空白字符占位
            color = self.TAGCOLOR[i % len(self.TAGCOLOR)]
            box = QCheckBox(tag)
            box.setStyleSheet(f"QCheckBox{{color:{color}; font-weight:bold}}")
            box.clicked.connect(
                lambda s, tag=tag: self.tag_selected_files(tag, s))
            self.tag_check_box_container.addItem(item)
            self.tag_check_box_container.setItemWidget(item, box)

    def add_tag_item(self):
        '''
        添加新的标签分类
        '''
        tag, flag = QInputDialog().getText(self, "添加标签", "添加标签",
                                           QLineEdit.EchoMode.Normal, "",
                                           Qt.WindowFlags(), Qt.InputMethodHints())
        if flag:
            self.tag_list.append(tag)
            self.fresh_tag_items()  # 更新标签选择框

    #--------------------------------------#
    #          图片展示窗口基础逻辑           #
    #--------------------------------------#
    def fresh_file_name_label(self):
        '''
        更新文件名上的类别图案
        '''
        for i in range(self.file_list_widget.count()):
            item = self.file_list_widget.item(i)
            file_name, tag_set = self.file_list[i + self.index_shift]
            # 获取文件标签在标签列表中的索引，用于标签对应的获取颜色
            tag_index = [self.tag_list.index(tag) for tag in tag_set]
            # 构造带彩色标签圆点的文件名字符串
            file_string = "".join([f"<font color={self.TAGCOLOR[i]}>●</font>" for i in tag_index if i < len(self.TAGCOLOR)])
            file_string += file_name
            # 设置label的显示内容
            label = self.file_list_widget.itemWidget(item)
            label.setText(QObject().tr(file_string))


    def fresh_page(self, page_num: int):
        '''
        更新当前图片展示页面
        '''
        # 判断目标页码合法性
        if page_num < 1 or page_num > self.page_num_total:
            return False
        # 更新成员变量和页数显示
        self.page_num = page_num
        self.index_shift = (self.page_num - 1) * self.files_per_page
        self.page_to_line_edit.setText(str(self.page_num))
        self.page_label.setText("/" + str(self.page_num_total))
        # 刷新图片
        self.file_list_widget.clear()
        for i in range(self.index_shift,
                       min(self.index_shift + self.files_per_page, self.file_count)):
            file_path = os.path.join(self.image_path, self.file_list[i][0])
            # 生成item对象
            item = QListWidgetItem(QIcon(file_path), "") # item本体，包含图片
            self.file_list_widget.addItem(item)
            label = QLabel() # 用于放置文件名字符串
            self.file_list_widget.setItemWidget(item, label)
        self.fresh_file_name_label()

    def zoom_request(self, a0: QWheelEvent) -> None:
        '''
        替换file_list_widget的滚轮事件
        增加Ctrl+滚轮的图片缩放功能
        '''
        if QApplication.keyboardModifiers() == Qt.ControlModifier:
            widget_size = self.file_list_widget.size()
            item_size = self.file_list_widget.iconSize()
            if a0.angleDelta().y() > 0 and (item_size.width()+12)*2 < widget_size.width():
                self.file_list_widget.setIconSize(
                    QSize(min(widget_size.width()//2-30, item_size.width()+50),
                          min(widget_size.height()-30, item_size.height())))
            elif a0.angleDelta().y() < 0 and (item_size.width()+12) * 5 > widget_size.width():
                self.file_list_widget.setIconSize(
                    QSize(max(widget_size.width()//5-12, item_size.width()-50),
                          item_size.height()))
            left_margin = widget_size.width() % (item_size.width() + 12) / (widget_size.width() // (item_size.width()+12))
            # self.file_list_widget.setContentsMargins(left_margin, 4, 4, 4)
            self.file_list_widget.setStyleSheet(f"QListWidgetItem{{margin-left:{left_margin}px;}}")
        else:
            return QListWidget.wheelEvent(self.file_list_widget, a0)

    def file_list_context_menu_event(self, event):
        '''
        Unfinished, context menu's checkbox don't work
        '''
        # rmenu = QMenu("标签")
        # for tag in self.tag_list:
        #     # action with out parent will be recoveried!
        #     tag_action = QAction(tag, rmenu, checkable=True)
        #     tag_action.setCheckable(True)
        #     rmenu.addAction(tag_action)
        # rmenu.mouseReleaseEvent = lambda e: None # 点击后不关闭菜单
        # rmenu.exec(QCursor.pos())
        pass

    #--------------------------------------#
    #             dock窗口基础逻辑            #
    #--------------------------------------#
    def fresh_statics_dock(self):
        '''
        标签文件个数统计窗口
        '''
        # 重建窗口速度比通过检索修改更高效
        NOTAG = "*未分类"
        # 清空窗口
        self.statics_tree.clear()
        # 统计变量, "标签": [数量， [文件1, 文件2...]]
        tag_statics = {tag:[0,[]] for tag in self.tag_list}
        tag_statics[NOTAG] = [0, []]
        # 遍历文件，统计所有文件的标签信息
        for file_name, tag_set in self.file_list:
            if not tag_set:
                tag_statics[NOTAG][0] += 1
                tag_statics[NOTAG][1].append(file_name) 
            for tag in tag_set:
                tag_statics[tag][0] += 1
                tag_statics[tag][1].append(file_name) 
        # 遍历标签，在TreeWidget中生成显示
        for tag, info in tag_statics.items():
            item = QTreeWidgetItem([tag])  # 初始化一级项目
            item.setData(1, Qt.DisplayRole, info[0]) # 修改计数显示，使用setData添加的数值可以被正确排序
            self.statics_tree.addTopLevelItem(item)  # 加入控件中
            # 初始化二级项目
            for i, file_name in enumerate(info[1]):
                child_item = QTreeWidgetItem()
                child_item.setData(0, Qt.DisplayRole, i+1)
                child_item.setData(1, Qt.DisplayRole, file_name)
                item.addChild(child_item)  # 增加下级项目

    #--------------------------------------#
    #             独立窗口基础逻辑            #
    #--------------------------------------#
    def show_single_image(self, file_index):
        '''
        独立窗口显示单个图像
        '''
        # 生成独立窗口
        if 'single_image_win' not in self.__dict__ \
                or self.single_image_win.isHidden():
            # 如果独立窗口处于关闭状态，则重新建立
            # 否则仅刷新图片，避免独立窗口占用焦点
            self.single_image_win = SingleImageWindow()
            self.single_image_win.show()
            # 切换图像

            def switch_image(directory: int):
                '''
                切换上一张或下一张图片
                directory=-1为上一张，directory=1为下一张
                '''
                # 更新index
                nonlocal file_index
                tmp_file_index = file_index + directory  # 获取在文件列表中的索引
                if tmp_file_index < 0 or tmp_file_index == self.file_count:
                    # 第一张图片和最后一张图片则不切换
                    return
                file_index = tmp_file_index
                item_index = file_index % self.files_per_page  # 获取当前页中的索引
                if directory == -1 and item_index == self.files_per_page - 1:
                    # 如果从该页第一张图片向前切换，则向前翻页
                    self.fresh_page(self.page_num - 1)
                elif directory == 1 and item_index == 0:
                    # 如果从该页最后一张图片向后切换，则向后翻页
                    self.fresh_page(self.page_num + 1)
                # 更新主窗口被选中的图片
                self.file_list_widget.clearSelection()
                self.file_list_widget.item(item_index).setSelected(True)
                # 更新独立窗口的图片
                file_name = self.file_list[file_index][0]
                file_path = os.path.join(self.image_path, file_name)
                self.single_image_win.set_image(file_path)
            QShortcut("a", self.single_image_win, lambda: switch_image(-1))
            QShortcut("d", self.single_image_win, lambda: switch_image(1))
        # 设定图像
        file_name = self.file_list[file_index][0]
        file_path = os.path.join(self.image_path, file_name)
        self.single_image_win.set_image(file_path)
        # self.single_image_win.center()

    #--------------------------------------#
    #              标签功能实现              #
    #--------------------------------------#
    def set_tag_checkbox(self, index: int):
        '''
        转换指定索引的checkbox状态，同时对已选文件设置该标签
        '''
        if index >= self.tag_check_box_container.count():
            return False
        item = self.tag_check_box_container.item(index)
        box = self.tag_check_box_container.itemWidget(item)
        box.setChecked(not box.isChecked())
        self.tag_selected_files(box.text(), box.isChecked())

    def tag_selected_files(self, tag: str, flag: bool):
        '''
        设定或移除已选文件标签
        '''
        # 将所选图片的listView索引转化为file_list中的索引
        selected_index = self.file_list_widget.selectedIndexes()
        if not selected_index:
            return
        index = [self.index_shift+si.row() for si in selected_index]
        # 修改file_list中的tag
        fun = set.add if flag else set.remove
        for i in index:
            fun(self.file_list[i][1], tag)
        # 更新统计信息
        self.fresh_file_name_label() # 刷新文件名上的标签圆点
        self.fresh_statics_dock()
        self.save_files_tag()

    def get_files_tag(self, *args):
        '''
        替换file_list_widget的selectionChanged事件
        增加根据已选文件标签，更新checkbox的功能
        '''
        QListWidget.selectionChanged(self.file_list_widget, *args)
        # -----清空所有复选框-----
        n = self.tag_check_box_container.count()
        for i in range(n):
            item = self.tag_check_box_container.item(i)
            box = self.tag_check_box_container.itemWidget(item)
            box.setChecked(False)

        # -----获取所选文件标签-----
        # 将所选图片的listView索引转化为file_list中的索引
        selected_index = self.file_list_widget.selectedIndexes()
        if not selected_index:
            return
        index = [self.index_shift+si.row() for si in selected_index]
        # 获取共有标签
        tags_set = self.file_list[index[0]][1]
        for i in index:
            tags_set &= self.file_list[i][1]  # 进行交集运算，获取共有标签

        # -----更新checkbox-----
        # 选中所选文件的共有标签
        for tag in tags_set:
            tag_index = self.tag_list.index(tag)
            item = self.tag_check_box_container.item(tag_index)
            box = self.tag_check_box_container.itemWidget(item)
            box.setChecked(True)

    #--------------------------------------#
    #               筛选部分部分             #
    #--------------------------------------#

    def remove_duplicate(self, tags, *, progresser=None):
        '''
        去重
        HACK 去重界面需要完善，增加进度条，增加操作确认，审阅将要去重的文件（较复杂），去重文件列表通过内存传递
        '''
        if not tags:
            return
        # 获取选定类别的文件
        file_list = []
        for file_name, tag_set in self.file_list:
            if not tag_set and "未分类" in tags :
                file_list.append(os.path.join(self.image_path, file_name))
                continue
            for tag in tag_set:
                if tag in tags:
                    file_list.append(os.path.join(self.image_path, file_name))
                    break
        # 去重处理
        hash_parser = DHash()
        hash_parser.find_duplicates_to_remove(file_list, outfile="dup.txt", progresser=progresser) # file为完整路径
        # 移动到临时文件夹
        if progresser:
            prgs_label, prgs_bar = progresser
            prgs_label.setText("3/3")
        with open("dup.txt", "r", encoding="utf-8") as f:
            content = f.read()
        file_list_to_remove = eval(content) # NOTE eval可能导致潜在的安全问题！！
        OUT_DIR = os.path.join(self.image_path, "dup")
        if not os.path.exists(OUT_DIR):
            os.mkdir(OUT_DIR)
        if progresser:
            prgs_bar.setRange(0, len(file_list_to_remove))
        for i, file_from in enumerate(file_list_to_remove):
            if progresser: prgs_bar.setValue(i+1)
            # 移动文件
            file_name = os.path.split(file_from)[1] # file_from为完整路径，提取其中的文件名
            file_to = os.path.join(OUT_DIR, file_name)
            shutil.move(file_from, file_to)
        # 更新file_list数据
        file_list_to_remove = [os.path.split(file)[1] for file in file_list_to_remove]
        file_removed_info = os.path.join(self.image_path, "dup", "dup_tag.txt")
        with open(file_removed_info, "w", encoding="utf-8") as f:
            i = 0
            while True:
                file_name, tag_set = self.file_list[i]
                if file_name in file_list_to_remove:
                    self.file_list.pop(i)
                    line = " ".join([file_name] + list(tag_set))
                    f.write(line)
                    f.write("\n")
                else:
                    i += 1
                    if i >= len(self.file_list):
                        break
        self.save_files_tag()
        print("完成！")
            
    def remove_duplicate_window(self):
        '''
        自动筛选窗口
        '''
        remove_duplicate_window = QWidget()
        # 标签选择
        checkbox_layout = QGridLayout()
        checkbox_list = []
        for i, tag in enumerate(self.tag_list+["未分类"]):
            checkbox = QCheckBox(tag)
            checkbox_layout.addWidget(checkbox, i//5, i % 5)
            checkbox_list.append(checkbox)
        # 进度显示
        prgs_label = QLabel()
        prgs_bar = QProgressBar()
        prgs_layout = QHBoxLayout()
        prgs_layout.addWidget(prgs_label)
        prgs_layout.addWidget(prgs_bar)
        # 确认按钮
        ok_button = QPushButton("确认")
        ok_button.clicked.connect(lambda: self.remove_duplicate(
            [box.text() for box in checkbox_list if box.isChecked()],
            progresser = [prgs_label, prgs_bar]
        ))
        cancel_button = QPushButton("取消")
        cancel_button.clicked.connect(remove_duplicate_window.close)
        button_layout = QHBoxLayout()
        # button_layout.addSpacing(500)
        button_layout.addWidget(ok_button)
        button_layout.addWidget(cancel_button)
        # 
        hbox = QHBoxLayout()
        hbox.addLayout(prgs_layout)
        hbox.addLayout(button_layout)
        # 主窗口
        main_layout = QVBoxLayout()
        main_layout.addLayout(checkbox_layout)
        main_layout.addLayout(hbox)
        remove_duplicate_window.setLayout(main_layout)
        remove_duplicate_window.setWindowTitle("去除近似")
        remove_duplicate_window.show()
        self.remove_duplicate_window = remove_duplicate_window

    #--------------------------------------#
    #               导入导出部分             #
    #--------------------------------------#
    def open_folder(self):
        '''
        选择目录，读取目录下图片和保存的标签信息
        '''
        image_path = QFileDialog().getExistingDirectory(self, "打开", os.getcwd())
        if not image_path:
            return
        self.image_path = os.path.normpath(image_path)
        self.tag_file = os.path.join(self.image_path, "tag.txt")
        self.tag_list = []
        self.file_list = [[f, set()] for f in os.listdir(self.image_path)
                          if f.split(".")[-1].lower() in self.IMAGETYPE]
        self.file_count = len(self.file_list)
        self.page_num_total = (self.file_count + self.files_per_page - 1) \
            // self.files_per_page

        # 载入初始页数据
        self.load_files_tag()
        self.fresh_tag_items()
        self.fresh_page(1)
        self.fresh_statics_dock()

    def export_files_by_tag(self, tags, export_path, *, progresser=None):
        '''
        导出包含指定标签的文件到指定目录下
        '''
        if not os.path.exists(export_path):
            os.mkdir(export_path)
        plabel, pbar = progresser
        with open(os.path.join(export_path, "log.txt"), 'w', encoding='utf-8') as f:
            plabel.setText(f"0/{len(tags)}")
            for i, tag in enumerate(tags):
                plabel.setText(f"{i+1}/{len(tags)}")
                file_list = [file[0]
                             for file in self.file_list if tag in file[1]]
                if not file_list:
                    continue
                export_tag_path = os.path.join(export_path, tag)
                if not os.path.exists(export_tag_path):
                    os.mkdir(export_tag_path)
                pbar.setRange(0, len(file_list))
                for j, file in enumerate(file_list):
                    pbar.setValue(j+1)
                    file_from = os.path.join(self.image_path, file)
                    file_to = os.path.join(export_tag_path, file)
                    shutil.copy(file_from, file_to)
                    f.write(f"copy {file_from}->{file_to}")
                    f.write("\n")

    def export_window(self):
        '''
        导出图片窗口
        '''
        export_window = QWidget()
        # 标签选择
        checkbox_layout = QGridLayout()
        checkbox_list = []
        for i, tag in enumerate(self.tag_list):
            checkbox = QCheckBox(tag)
            checkbox_layout.addWidget(checkbox, i//5, i % 5)
            checkbox_list.append(checkbox)
        # 导出路径设置
        label = QLabel("导出到：")
        path_edit = QLineEdit()
        path_setting_layout = QHBoxLayout()
        path_setting_layout.addWidget(label)
        path_setting_layout.addWidget(path_edit)
        # 进度显示
        prgs_label = QLabel()
        prgs_bar = QProgressBar()
        prgs_layout = QHBoxLayout()
        prgs_layout.addWidget(prgs_label)
        prgs_layout.addWidget(prgs_bar)
        # 确认按钮
        ok_button = QPushButton("确认")
        ok_button.clicked.connect(lambda: self.export_files_by_tag(
            [box.text() for box in checkbox_list if box.isChecked()],
            path_edit.text(), progresser=[prgs_label, prgs_bar]
        ))
        cancel_button = QPushButton("取消")
        cancel_button.clicked.connect(export_window.close)
        button_layout = QHBoxLayout()
        # button_layout.addSpacing(500)
        button_layout.addWidget(ok_button)
        button_layout.addWidget(cancel_button)
        # 
        hbox = QHBoxLayout()
        hbox.addLayout(prgs_layout)
        hbox.addLayout(button_layout)
        # 主窗口
        main_layout = QVBoxLayout()
        main_layout.addLayout(checkbox_layout)
        main_layout.addLayout(path_setting_layout)
        main_layout.addLayout(hbox)
        export_window.setLayout(main_layout)
        export_window.setWindowTitle("导出分类")
        export_window.show()
        self.export_window = export_window

    def load_files_tag(self):
        '''
        读取已保存的标签列表和文件标签
        '''
        if not os.path.exists(self.tag_file):
            return  # 如果不存在已经保存的标签信息则不读取
        with open(self.tag_file, 'r', encoding='utf-8') as f:
            # 读取标签列表
            tags_str = f.readline()
            self.tag_list = tags_str.strip().split(" ")
            # 读取文件标签信息
            file_list = [item[0] for item in self.file_list]
            # TODO 严重减慢加载速度！
            # TODO 处理标签文件不存在的情况
            for line in f.readlines():
                file_name, *tags = line.split()
                idx = file_list.index(file_name)
                if idx != -1:
                    self.file_list[idx][1] = set(tags)

    def save_files_tag(self):
        '''
        保存标签列表和文件标签
        '''
        #
        if not self.image_path:
            return  # 如果没有打开目录则不保存
        with open(self.tag_file, 'w', encoding='utf-8') as f:
            # 写标签列表
            f.write(" ".join(self.tag_list))
            f.write("\n")
            # 写文件标签信息
            for file_tag in self.file_list:
                line_str = file_tag[0] + " "
                line_str += " ".join(sorted(file_tag[1]))
                line_str += "\n"
                f.write(line_str)


class SingleImageWindow(QWidget):
    def __init__(self):
        super().__init__()
        # 显示图片
        self.label_image = QLabel(self)
        layout = QVBoxLayout()
        layout.addWidget(self.label_image)
        self.setLayout(layout)
        # 快捷键
        QShortcut("Esc", self, self.close)
        QShortcut("Ctrl+w", self, self.close)

    def center(self):
        # 居中显示
        # 由于要加载图片，窗口显示之后，才能得到正确尺寸
        # 所以需要在super().show()之后执行居中
        screen = QDesktopWidget().screenGeometry()
        size = self.geometry()
        self.move((screen.width() - size.width()) // 2,
                  (screen.height() - size.height()) // 2)

    def resize_by_image(self):
        '''
        根据image的尺寸调整窗口大小
        '''
        size = self.label_image.size()
        self.resize(size)

    def set_image(self, file_path):
        '''
        显示指定文件
        '''
        self.setWindowTitle(os.path.normpath(file_path))
        image = QPixmap(file_path)
        self.label_image.setPixmap(image)
        self.resize_by_image()

if __name__ == "__main__":
    app = QApplication(sys.argv)
    win = WindowClass()
    win.show()
    sys.exit(app.exec_())