import os
from pathlib import Path
import cv2
import numpy as np
from ultralytics import YOLO
import shutil

# ---------------------------
# 全局配置
# ---------------------------

# ✅ 定义唯一的 ROI (x, y, w, h) —— 请根据你的实际坐标修改！
SINGLE_ROI = (859,810,696,328)  # 格式: x, y, w, h

CLASS_NAMES = {
    0: "未堆料",
    1: "小堆料",
    2: "大堆料",
    3: "未浇筑满",
    4: "浇筑满"
}

TARGET_SIZE = 640
THRESHOLD = 0.4  # 加权得分阈值


# ---------------------------
# class1/class2 加权判断
# ---------------------------
def weighted_small_large(pred_probs, threshold=0.4, w1=0.3, w2=0.7):
    p1 = float(pred_probs[1])
    p2 = float(pred_probs[2])
    total = p1 + p2
    if total > 0:
        score = (w1 * p1 + w2 * p2) / total
    else:
        score = 0.0
    final_class = "大堆料" if score >= threshold else "小堆料"
    return final_class, score, p1, p2


# ---------------------------
# 批量推理主函数（单 ROI + 批量处理）
# ---------------------------
def batch_classify_images(model_path, input_folder, output_root, target_size=640, threshold=0.4):
    model = YOLO(model_path)

    output_root = Path(output_root)
    output_root.mkdir(parents=True, exist_ok=True)

    # 创建每个类别的输出目录
    class_dirs = {}
    for name in CLASS_NAMES.values():
        d = output_root / name
        d.mkdir(exist_ok=True)
        class_dirs[name] = d

    x, y, w, h = SINGLE_ROI

    # 存储所有裁剪后的图像和对应的原始路径
    crops = []
    img_paths = []

    input_folder = Path(input_folder)
    supported_exts = {'.jpg', '.jpeg', '.png', '.bmp', '.tif'}

    print("🔍 正在裁剪所有图片的 ROI 区域...")
    for img_path in input_folder.glob("*.*"):
        if img_path.suffix.lower() not in supported_exts:
            continue
        img = cv2.imread(str(img_path))
        if img is None:
            continue

        # 裁剪并 resize 到模型输入大小
        roi = img[y:y+h, x:x+w]
        roi_resized = cv2.resize(roi, (target_size, target_size), interpolation=cv2.INTER_AREA)
        crops.append(roi_resized)
        img_paths.append(img_path)

    if not crops:
        print("❌ 没有有效图片可供处理")
        return

    total = len(crops)
    print(f"✅ 共准备 {total} 张图片，开始批量推理...")

    # 🔥 批量推理（一次 forward pass）
    results = model(
        source=crops,
        verbose=False,
        imgsz=target_size,
        half=False,      # 如果使用 GPU 可开启半精度: half=True
        device=0         # 使用 GPU 0，如用 CPU 改为: device=None
    )

    # 后处理结果
    print("📦 正在处理结果并保存...")
    for i, result in enumerate(results):
        pred_probs = result.probs.data.cpu().numpy().flatten()
        class_id = int(pred_probs.argmax())
        confidence = float(pred_probs[class_id])
        class_name = CLASS_NAMES.get(class_id, f"未知类别({class_id})")

        if class_id in [1, 2]:
            final_class, score, p1, p2 = weighted_small_large(pred_probs, threshold=threshold)
        else:
            final_class = class_name
            score = confidence
            p1 = float(pred_probs[1])
            p2 = float(pred_probs[2])

        # 构造目标路径
        original_path = img_paths[i]
        suffix = f"_roi0_{final_class}_score{score:.2f}_p1{p1:.2f}_p2{p2:.2f}"
        dst_path = class_dirs[final_class] / f"{original_path.stem}{suffix}{original_path.suffix}"

        # 复制原图到对应类别文件夹
        shutil.copy2(str(original_path), str(dst_path))
        print(f"{original_path.name} -> {final_class} (score={score:.2f}, p1={p1:.2f}, p2={p2:.2f})")

    print(f"\n🎉 分类完成！共处理 {total} 张图片")


# ---------------------------
# 使用示例
# ---------------------------
if __name__ == "__main__":
    MODEL_PATH = r"cls5.pt"
    INPUT_FOLDER = r"./test_image"
    OUTPUT_ROOT = r"./classified_images"

    batch_classify_images(MODEL_PATH, INPUT_FOLDER, OUTPUT_ROOT, TARGET_SIZE, THRESHOLD)