import time
import fastdeploy as fd
import cv2


class YoloV8Prediction:
    def __init__(self, ModelPath: str, PreheatingImagesPath: str):
        option = fd.RuntimeOption()
        # 切换使用CPU/GPU
        # option.use_cpu()
        option.use_gpu()
        # 切换不同后端
        # option.use_paddle_infer_backend()
        # option.use_paddle_lite_backend()
        # option.use_paddle_backend()
        option.use_openvino_backend()
        # option.use_ort_backend()
        option.use_trt_backend()
        # option.use_poros_backend()
        self.yolov8 = fd.vision.detection.YOLOv8(ModelPath, runtime_option=option)
        self.yolov8.preprocessor.size = [640, 640]
        img = cv2.imread(PreheatingImagesPath)
        s = time.time()
        for i in range(100):
            self.yolov8.predict(img)
        print(f'模型预热成功，预测100张图片耗时{time.time() - s}秒')

    def Start_Prediction(self, Images: str):
        return self.yolov8.predict(Images)

from ultralytics import YOLO
import torch
import glob
import os
import xml.etree.ElementTree as ET
import random
import shutil
import time
from window.windowclass_en import WindowManager
import cv2
import Draw.D3Gui
from yolov8.YoloV8_Prediction import YoloV8Prediction


class YoloV8:
    def __init__(self, ModelPath: str, ParentWindowHandle, WindowClassName: str, WindowTitle: str,
                 PreheatingImagesPath: str, Mode: int = 1):
        """
        使用YOLOV8时常用的函数，mode为0时为训练模型，为1时为预测模式（训练和预测的环境不一样，训练支支持64位python。训练是cuda+torch+安装ultralytics，预测是TensorRT+fastdeploy）
        :param ModelPath: str,模型路径，训练模式为pt模型，预测模式为onnx模型
        :param ParentWindowHandle: int，父窗口句柄
        :param WindowClassName: str，窗口类名
        :param WindowTitle: str，窗口标题
        :param PreheatingImagesPath: str，用于预热预测模型的图片路径
        :param Mode: int，运行模式（训练/预测）
        """
        self.mode = Mode
        self.window = WindowManager()
        self.handle = self.window.Find_Windows(ParentWindowHandle, WindowClassName, WindowTitle)
        if Mode == 0:
            self.yolo = YOLO(model=ModelPath)
        elif Mode == 1:
            self.yolo = YoloV8Prediction(ModelPath, PreheatingImagesPath)
        pass

    @staticmethod
    # 转换一个xml文件为txt
    def single_xml_to_txt(xml_file, class_names):
        tree = ET.parse(xml_file)
        root = tree.getroot()
        # 保存txt文件路径
        text_dir, xml_name = os.path.split(xml_file)
        txt_name = f'{xml_name[:-4]}.txt'
        txt_file = os.path.join(text_dir, txt_name)
        with open(txt_file, 'w') as txt_file:
            for member in root.findall('object'):
                # 从xml获取图像的宽和高
                picture_width = int(root.find('size')[0].text)
                picture_height = int(root.find('size')[1].text)
                class_name = member[0].text

                # 类名对应的index
                class_num = class_names.index(class_name)
                box_x_min = int(member[4][0].text)  # 左上角横坐标
                box_y_min = int(member[4][1].text)  # 左上角纵坐标
                box_x_max = int(member[4][2].text)  # 右下角横坐标
                box_y_max = int(member[4][3].text)  # 右下角纵坐标
                # 转成相对位置和宽高（所有值处于0~1之间）
                x_center = (box_x_min + box_x_max) / (2 * picture_width)
                y_center = (box_y_min + box_y_max) / (2 * picture_height)
                width = (box_x_max - box_x_min) / picture_width
                height = (box_y_max - box_y_min) / picture_height
                print(class_num, x_center, y_center, width, height)
                txt_file.write(
                    str(class_num) + ' ' + str(x_center) + ' ' + str(y_center) + ' ' + str(width) + ' ' + str(
                        height) + '\n')

    @staticmethod
    def Batch_Rename(dir: str, ExtensionName: str):
        """
        批量重命名
        :param dir: str,文件目录
        :param ExtensionName: str,文件拓展名
        :return:
        """
        files = os.listdir(dir)
        i = 1
        for file in files:
            if ExtensionName in file:
                new_name = f"{i}.{ExtensionName}"
                i += 1
                old_path = os.path.join(dir, file)
                new_path = os.path.join(dir, new_name)
                os.rename(old_path, new_path)

    @staticmethod
    def Delete_obsolete_images(image_folder: str, label_folder: str):
        """
        匹配txt和jpg的数量,删除废图
        :param image_folder: str，图集路径
        :param label_folder: str,文本集路径
        :return:
        """
        # 获取label文件夹中的所有.txt文件
        label_files = os.listdir(label_folder)
        # 构建一个集合来存储存在于label中的文件名（不包括扩展名）
        label_file_names = {os.path.splitext(file)[0] for file in label_files}
        # 获取image文件夹中的所有.jpg文件
        image_files = os.listdir(image_folder)
        # 遍历每个.jpg文件
        for image_file in image_files:
            # 提取文件名（不包括扩展名）
            image_file_name = os.path.splitext(image_file)[0]
            # 检查文件名是否存在于label中
            if image_file_name not in label_file_names:
                # 如果文件名不在label中，输出提示信息并删除该.jpg文件
                print(f"删除多余的图像文件: {os.path.join(image_folder, image_file)}")
                os.remove(os.path.join(image_folder, image_file))

    @staticmethod
    def Data_Partitioning(image_folder, label_folder):
        """
        按照7：2：1划分train：val：test，不会删除源文件
        :param image_folder:  str,图集路径
        :param label_folder:  str,标注集路径
        :return:
        """
        # 创建train、val和test子文件夹
        os.makedirs(os.path.join(image_folder, 'train'), exist_ok=True)
        os.makedirs(os.path.join(image_folder, 'val'), exist_ok=True)
        os.makedirs(os.path.join(image_folder, 'test'), exist_ok=True)
        os.makedirs(os.path.join(label_folder, 'train'), exist_ok=True)
        os.makedirs(os.path.join(label_folder, 'val'), exist_ok=True)
        os.makedirs(os.path.join(label_folder, 'test'), exist_ok=True)
        # 获取image文件夹中的所有.jpg文件和label文件夹中的所有.txt文件
        image_files = 
        label_files = 
        # 随机打乱文件列表
        random.shuffle(image_files)
        random.shuffle(label_files)
        # 计算划分的数量
        total_files = len(image_files)
        train_ratio = 0.7
        val_ratio = 0.2
        test_ratio = 0.1
        train_count = int(train_ratio * total_files)
        val_count = int(val_ratio * total_files)
        # 将文件拷贝到对应的子文件夹中
        for i, image_file in enumerate(image_files):
            if i < train_count:
                dest_folder = 'train'
            elif i < train_count + val_count:
                dest_folder = 'val'
            else:
                dest_folder = 'test'
            # 拷贝图像文件
            shutil.copy(os.path.join(image_folder, image_file), os.path.join(image_folder, dest_folder, image_file))
            # 构建对应的label文件名
            label_file = image_file.replace('.jpg', '.txt')
            # 拷贝label文件
            shutil.copy(os.path.join(label_folder, label_file), os.path.join(label_folder, dest_folder, label_file))

    @staticmethod
    def Training_Model(ModelPath: str, DataPath: str, Epochs: int, Imgsz: int, Batch: int):
        """
        加载预先训练的模型用于训练，支持断点训练
        :param ModelPath: str,模型路径
        :param DataPath: str,数据集路径，eg:xxx/slide.yaml
        :param Epochs: int,训练轮数
        :param Imgsz: int,图像大小，eg：320，640
        :param Batch: int,显存大小，eg:2,4,6,8
        :return:
        """
        model = YOLO(ModelPath)  # 根据YAML构建并传递权重
        model.train(data=DataPath, epochs=Epochs, imgsz=Imgsz, batch=Batch, cache=True)

    @staticmethod
    def Build_Weights_And_Train(WeightFilePath: str, ModelPath: str, DataPath: str, Epochs: int, Imgsz: int,
                                Batch: int):
        """
        根据YAML构建并传递权重，然后开始训练
        :param WeightFilePath: str,权重路径
        :param ModelPath: str,模型路径
        :param DataPath: str,数据集路径，eg:xxx/slide.yaml
        :param Epochs: int,训练轮数
        :param Imgsz: int,图像大小，eg：320，640
        :param Batch: int,显存大小，eg:2,4,6,8
        :return:
        """
        model = YOLO(WeightFilePath).load(ModelPath)  # 根据YAML构建并传递权重
        model.train(data=DataPath, epochs=Epochs, imgsz=Imgsz, batch=Batch, cache=True)

    def Dir_Xml_To_Txt(self, path: str, class_names):
        """
        将xml标注转换为yolo标注
        :param path: str,标注集路径, 后面的斜杠要带
        :param class_names: 类名，eg：class_names = ['fei', 'jing']
        :return:
        """
        #  转换文件夹下的所有xml文件为txt
        for xml_file in glob.glob(path + '*.xml'):  # 查找符合规定的文件路径名
            print(xml_file)
            self.single_xml_to_txt(xml_file, class_names)

    def Using_CPU_Prediction(self, ImagePath: str):
        """
        使用CPU预测(仅用于训练时的预测测试)
        :param ImagePath: str,要预测的图片路径
        :return:
        """
        if self.mode == 0:
            self.yolo.predict(ImagePath, save=True)
        else:
            print("为保证性能预测模式不支持调用该函数")

    def Using_GPU_Prediction(self, ImagePath: str):
        """
        使用GPU预测(仅用于训练时的预测测试)
        :param ImagePath: str,要预测的图片路径
        :return:
        """
        if self.mode == 0:
            torch.cuda.is_available()
            self.yolo.predict(ImagePath, device="0")
        else:
            print("为保证性能预测模式不支持调用该函数")

    def Screenshot_Always(self, Path: str):
        """
        窗口截图，一直截（如果黑屏请更换模式Capture_Win32_Gpu（handle,mode））
        :param Path: str，图片保存路径，最后不要带/
        :return:
        """
        if not self.handle:
            print("获取窗口句柄失败")
            exit()
        i = 0
        try:
            while True:  # 无限循环
                image = self.window.Capture_Win32_Gpu(self.handle[0], 3)
                cv2.imwrite(f"{Path}/{i}.jpg", image)
                i += 1
                time.sleep(1)
        except KeyboardInterrupt:
            print("用户手动停止截图")

    def Screenshot_Breakpoint(self, start: int, end: int, Path: str):
        """
        窗口截图，指定数量（如果黑屏请更换模式Capture_Win32_Gpu（handle,mode））
        :param start: int，从几开始
        :param end: int，到几结束
        :param Path: str，图片保存路径，最后不要带/
        :return:
        """
        if not self.handle:
            print("获取窗口句柄失败")
            exit()
        else:
            for i in range(start, end + 1):
                image = self.window.Capture_Win32_Gpu(self.handle[0], 3)
                cv2.imwrite(f"{Path}/{i}.jpg", image)
                time.sleep(1)

    def Draw_Rectangle(self):
        """
        矩形绘制
        :return:
        """
        if self.mode == 0:
            if not self.handle:
                print("获取窗口句柄失败")
                exit()
            draw = Draw.D3Gui.ExecDraw(self.handle[0])  # 初始化模块
            while True:
                image = self.window.Capture_Win32_Gpu(self.handle[0], 3)
                result = self.yolo.Start_Prediction(image)
                if result:
                    for box in result.boxes:
                        # 获取四个顶点坐标
                        left = int(box[0])
                        top = int(box[1])
                        right = int(box[2])
                        bottom = int(box[3])
                        draw.startLoop()  # 开始绘制段
                        draw.drawRect(left, top, right, bottom, 5, (255, 254, 0))  # 绘制矩形框
                        draw.endLoop()  # 结束绘制段
        else:
            print("为保证性能预测模式不支持调用该函数")

    def Automatic_Dimension(self, ImagesPath: str, LabelsPath: str, ImageSuffix: str):
        """
        自动标注
        :return:
        """
        if self.mode == 1:
            image_files = 

            for file_name in image_files:
                # 分割文件名和后缀
                parts = file_name.split('.')
                # 检查是否有两个部分，并且后缀是指定后缀
                if len(parts) == 2 and parts[1] == ImageSuffix.replace(".", ""):
                    imagepath = f'{ImagesPath}/{file_name}'
                    img = cv2.imread(imagepath)
                    img_height, img_width, _ = img.shape
                    result = self.yolo.Start_Prediction(img)

                    if result:
                        with open(f'{LabelsPath}/{parts[0]}.txt', "a") as file:
                            for a in range(len(result.label_ids)):
                                c = result.label_ids[a]
                                x_min = result.boxes[a][0]
                                y_min = result.boxes[a][1]
                                x_max = result.boxes[a][2]
                                y_max = result.boxes[a][3]

                                # 计算归一化坐标
                                x_center = (x_min + x_max) / (2 * img_width)
                                y_center = (y_min + y_max) / (2 * img_height)
                                width = (x_max - x_min) / img_width
                                height = (y_max - y_min) / img_height

                                # 将归一化坐标写入文件
                                normalized_box_str = f'{c} {x_center} {y_center} {width} {height}\n'
                                file.write(normalized_box_str)
        else:
            print("为保证性能预测模式不支持调用该函数")


if __name__ == '__main__':
    yolo = YoloV8('models/csgo_0.onnx', None, 'Valve001', 'Counter-Strike', 'CSGO.jpg', 1)
    yolo.Automatic_Dimension(r'C:\Users\devcat\Desktop\CSGO\images\train', r'C:\Users\devcat\Desktop\瓦\datasets'
                                                                           r'\slide\YOLO\labels', '.jpg')