import cv2
from ultralytics import YOLO
import numpy as np
from PyQt5.QtWidgets import QApplication, QFileDialog
import os
import shutil
from datetime import datetime


class yolo_class():
    def __init__(self, yolo_path):
        self.yolo_path = yolo_path
        self.model = YOLO(yolo_path)  # load a custom model

    def YoloDetect(self, imgs):
        results = self.model(imgs)  # predict on an image
        return results


def open_image_files():
    # 打開文件對話框，設置默認路徑
    file_dialog = QFileDialog()
    file_dialog.setDirectory(r"D:\Screwdriver_image\AOI_Origin_Error\Error")
    file_paths, _ = file_dialog.getOpenFileNames(None, "Select Images", "", "Image Files (*.png *.jpg *.bmp)")
    return file_paths


def get_class_name(index):
    # 更新類型名稱對應表，包含 41 個類別
    class_names = {
        0: "HexTamperproof_1", 1: "HexTamperproof_2", 2: "HexTamperproof_3", 3: "HexTamperproof_4",
        4: "Hex_1", 5: "Hex_2", 6: "Hex_3",
        7: "Opposite_1",
        8: "Pentalope_1", 9: "Pentalope_2",
        10: "Phillips_1", 11: "Phillips_2", 12: "Phillips_3", 13: "Phillips_4", 14: "Phillips_5",
        15: "Slotted_1", 16: "Slotted_2", 17: "Slotted_3", 18: "Slotted_4", 19: "Slotted_5", 20: "Slotted_6",
        21: "Spanner_1", 22: "Spanner_2",
        23: "Square_1", 24: "Square_2",
        25: "Standoff_1",
        26: "TORXTamperproof_1", 27: "TORXTamperproof_2", 28: "TORXTamperproof_3", 29: "TORXTamperproof_4",
        30: "TORXTamperproof_5", 31: "TORXTamperproof_6", 32: "TORXTamperproof_7",
        33: "TORX_1", 34: "TORX_2",
        35: "TriWing_1", 36: "TriWing_2", 37: "TriWing_3",
        38: "Triangle_1", 39: "Triangle_2",
        40: "Vacancy_1"
    }
    return class_names.get(index, f"Unknown_{index}")

        # 0: "phillips", 1: "triwing", 2: "standoff", 3: "spanner", 4: "hex",
        # 5: "hextamperproof", 6: "slotted", 7: "torx", 8: "torxtamperproof",
        # 9: "square", 10: "triangle", 11: "pentalope", 12: "phillips_1",
        # 13: "phillips_2", 14: "phillips_3", 15: "phillips_4", 16: "vacancy",
        # 17: "opposite"

        # 0: "HexTamperproof_1", 1: "HexTamperproof_2", 2: "HexTamperproof_3", 3: "HexTamperproof_4",
        # 4: "Hex_1", 5: "Hex_2", 6: "Hex_3",
        # 7: "Opposite_1",
        # 8: "Pentalope_1", 9: "Pentalope_2",
        # 10: "Phillips_1", 11: "Phillips_2", 12: "Phillips_3", 13: "Phillips_4", 14: "Phillips_5",
        # 15: "Slotted_1", 16: "Slotted_2", 17: "Slotted_3", 18: "Slotted_4", 19: "Slotted_5", 20: "Slotted_6",
        # 21: "Spanner_1", 22: "Spanner_2",
        # 23: "Square_1", 24: "Square_2",
        # 25: "Standoff_1",
        # 26: "TORXTamperproof_1", 27: "TORXTamperproof_2", 28: "TORXTamperproof_3", 29: "TORXTamperproof_4",
        # 30: "TORXTamperproof_5", 31: "TORXTamperproof_6", 32: "TORXTamperproof_7",
        # 33: "TORX_1", 34: "TORX_2",
        # 35: "TriWing_1", 36: "TriWing_2", 37: "TriWing_3",
        # 38: "Triangle_1", 39: "Triangle_2",
        # 40: "Vacancy_1"
def process_and_rename_images(yolo, image_paths, base_output_dir):
    for image_path in image_paths:
        image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
        results = yolo.YoloDetect(image)

        result_Max_Index = results[0].probs.top1  # 獲取分類編號
        confidence = results[0].probs.top1conf.item()  # 獲取 top1 的置信度

        # 獲取類型名稱
        class_name = get_class_name(result_Max_Index)

        print(f"原始圖片: {image_path}")
        print(f"推論結果: {class_name}, 信心值: {confidence}")

        # 只有當信心值高於 0.6 時才保存圖片
        if confidence > 0.4:
            # 獲取當前日期和時間，精度到毫秒
            current_time = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-1]

            # 獲取原始檔案的副檔名
            _, file_extension = os.path.splitext(image_path)

            # 創建新的檔名，包含日期時間
            new_filename = f"{class_name}_conf{confidence:.8f}{file_extension}"

            # 創建對應類別的資料夾
            class_output_dir = os.path.join(base_output_dir, class_name)
            if not os.path.exists(class_output_dir):
                os.makedirs(class_output_dir)

            # 設置新檔案的完整路徑
            new_filepath = os.path.join(class_output_dir, new_filename)

            # 複製並重命名檔案
            shutil.copy2(image_path, new_filepath)

            print(f"新檔名: {new_filename}")
            print(f"儲存位置: {class_output_dir}")
            print("儲存成功")
        else:
            print("信心值過低，不儲存圖片")

        print("---")


if __name__ == '__main__':
    app = QApplication([])
    yolo = yolo_class(r"D:\ScrewdriverFile\train_0811_Milwaukee_41\weights\best.pt")
    # D:\ScrewdriverFile\train_0811_Milwaukee_41\weights\best.pt
    # C:\Users\user\Desktop\0723ProgramBackup\model\1c344897-1978-4664-9bbd-256be14a3125.pt
    image_paths = open_image_files()
    if image_paths:
        base_output_dir = r"D:\Screwdriver_image\AOI_Origin_Error\Error_rename"
        process_and_rename_images(yolo, image_paths, base_output_dir)
    else:
        print("未選擇任何圖片")