import os
import yaml
from ultralytics import YOLO


def create_yaml_config(dataset_path, yaml_output_path):
    """
    自动生成 YOLO 需要的 data.yaml 配置文件
    这里的类别 ID 和名称必须和前端、Django 视图里保持绝对一致！
    """
    data = {
        'path': r"D:\eggplantshapeBS1\djangoProject\media\evolve_dataset",  # 数据集根目录的绝对路径
        'train': 'images/train',  # 训练集图片的相对路径
        'val': 'images/train',  # 演进时样本较少，验证集直接复用训练集
        'names': {
            0: 'Fruit Rot',
            1: 'Fruit borer',
            2: 'Healthy',
            3: 'Melon Thrips'
        }
    }
    with open(yaml_output_path, 'w', encoding='utf-8') as f:
        yaml.dump(data, f, allow_unicode=True)
    return yaml_output_path


def start_yolo_evolution(current_weight_path, dataset_dir, output_dir):
    print("🚀 紫睛灵：开始模型极速自演进...")

    # 1. 检查是否存在困难样本
    images_dir = os.path.join(dataset_dir, 'images', 'train')
    if not os.path.exists(images_dir) or len(os.listdir(images_dir)) == 0:
        print("❌ 演进终止：暂无需要学习的困难样本。请先在网页端提交纠错反馈。")
        return None

    # 2. 自动生成专属的 yaml 配置文件
    abs_dataset_dir = os.path.abspath(dataset_dir)
    yaml_path = os.path.join(abs_dataset_dir, 'evolve_data.yaml')
    create_yaml_config(abs_dataset_dir, yaml_path)
    print(f"✅ 已自动生成数据集配置: {yaml_path}")
    print(f"✅ 发现困难样本，即将基于 {current_weight_path} 启动防遗忘微调...")

    # 3. 加载当前的 YOLO11n 权重
    model = YOLO(current_weight_path)

    # 4. 启动轻量化微调训练 (核心参数已调优)
    results = model.train(
        data=yaml_path,
        epochs=15,  # 训练 15 轮，足够让模型记住错题
        batch=8,  # 增量学习样本少，batch size 设小一点
        imgsz=640,
        lr0=1e-4,  # 极低的学习率！确保不发生“灾难性遗忘”
        freeze=10,  # 【核心保命参数】冻结前 10 层骨干网络，保证参数量绝对不涨！
        project=output_dir,
        name="evolution_latest",
        exist_ok=True,  # 允许覆盖之前的演进文件夹
        workers=2,
        save=True,
        device=0,
    )

    # 5. 输出演进结果
    new_weight_path = os.path.join(output_dir, "evolution_latest", "weights", "best.pt")
    print("\n" + "=" * 50)
    print(f"🎉 演进大功告成！紫睛灵已进化！")
    print(f"📦 新模型权重已生成于: {new_weight_path}")
    print("💡 下一步：把这个新的 best.pt 拷贝并替换掉你 Django 系统目前正在使用的权重即可！")
    print("=" * 50 + "\n")


if __name__ == "__main__":
    # =====================================================================
    # ⚠️ 运行前，你只需要修改以下 3 个路径参数！
    # =====================================================================

    # 1. 你当前正在使用的 YOLO11n 模型权重路径 (相对路径或绝对路径)
    # 比如: "weights/best.pt" 或者 "yolo11n.pt"
    CURRENT_MODEL = r"D:\eggplantshapeBS1\djangoProject\PAD_train\src\src\results\exp4\weights\best.pt"

    # 2. Django 保存困难样本的根目录 (必须是指向包含 images 和 labels 的那个文件夹)
    # 根据我们刚才 views.py 的逻辑，它在 media 目录下
    DATASET_DIR = "media/evolve_dataset"

    # 3. 演进后新模型的保存目录
    OUTPUT_DIR = "runs/evolve"

    # =====================================================================

    start_yolo_evolution(
        current_weight_path=CURRENT_MODEL,
        dataset_dir=DATASET_DIR,
        output_dir=OUTPUT_DIR
    )