"""Phase 1: 半监督学习 + 主动学习 — 立即可执行方案

目标: 利用 3000 张未标注图，提升 mAP 到 0.65+

Step 1: 用最佳模型生成伪标签
Step 2: 主动学习筛选高价值样本
Step 3: 人工快速审核
Step 4: 重新训练

预期时间: 2-3 天
预期提升: +5~10%
"""
import os
import sys
import numpy as np
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))


def step1_generate_pseudo_labels():
    """用最佳模型预测未标注数据，生成伪标签"""
    from ultralytics import YOLO

    # 加载最佳模型（明早从 overnight 实验中选）
    best_model_path = "runs/detect/Detection_experiments/ablation_D_cbam_p2_wiou/weights/best.pt"
    model = YOLO(best_model_path)

    # 预测未标注图片
    unlabeled_dir = "Data/Detection_dataset/images/unlabeled"  # 你的 3000 张未标注图
    results = model.predict(
        source=unlabeled_dir,
        save_txt=True,
        save_conf=True,
        conf=0.3,  # 低阈值，保留更多候选
        iou=0.6,
        imgsz=1280,
        project="pseudo_labels",
        name="round1",
    )

    print(f"✓ 生成伪标签完成: pseudo_labels/round1/labels/")
    return "pseudo_labels/round1/labels/"


def step2_active_learning_selection(pseudo_label_dir, top_k=500):
    """主动学习：选出最有价值的 top_k 张图让人标注

    不确定性估计方法:
    1. 置信度方差 (同一目标多次预测的置信度波动)
    2. 框数量异常 (过多或过少)
    3. 小目标占比高
    """
    import glob
    from collections import defaultdict

    label_files = glob.glob(os.path.join(pseudo_label_dir, "*.txt"))

    uncertainty_scores = []

    for lbl_file in label_files:
        with open(lbl_file) as f:
            lines = [l.strip().split() for l in f if l.strip()]

        if len(lines) == 0:
            # 空图，可能是漏检
            uncertainty_scores.append((lbl_file, 10.0, "empty"))
            continue

        # 提取置信度和框大小
        confs = [float(l[5]) for l in lines if len(l) >= 6]
        areas = [float(l[3]) * float(l[4]) for l in lines]

        # 不确定性指标
        conf_std = np.std(confs) if len(confs) > 1 else 0
        low_conf_ratio = sum(1 for c in confs if c < 0.5) / len(confs)
        small_obj_ratio = sum(1 for a in areas if a < 0.002) / len(areas)
        num_boxes = len(lines)

        # 综合得分（越高越不确定）
        score = (
            conf_std * 2.0 +
            low_conf_ratio * 3.0 +
            small_obj_ratio * 1.5 +
            abs(num_boxes - 31) / 31 * 1.0  # 偏离平均数
        )

        uncertainty_scores.append((lbl_file, score, f"conf_std={conf_std:.2f}"))

    # 排序，选 top_k
    uncertainty_scores.sort(key=lambda x: x[1], reverse=True)
    selected = uncertainty_scores[:top_k]

    # 保存选中的文件列表
    output_file = "active_learning_selected.txt"
    with open(output_file, "w") as f:
        for lbl_file, score, reason in selected:
            img_file = lbl_file.replace("/labels/", "/images/").replace(".txt", ".jpg")
            f.write(f"{img_file}\t{score:.3f}\t{reason}\n")

    print(f"✓ 主动学习选择完成: {top_k} 张高价值样本")
    print(f"  保存到: {output_file}")
    print(f"  平均不确定性: {np.mean([s[1] for s in selected]):.3f}")

    return output_file


def step3_human_review_tool(selected_file):
    """快速人工审核工具

    显示图片 + 伪标签，人工确认/修改/删除
    """
    print("\n=== 人工审核工具 ===")
    print("使用 LabelImg 或 CVAT 快速审核:")
    print(f"  1. 打开 {selected_file}")
    print(f"  2. 逐张检查伪标签")
    print(f"  3. 修改错误的框")
    print(f"  4. 删除误检")
    print(f"  5. 补充漏检")
    print()
    print("预计时间: 500 张 × 30秒/张 = 4 小时")
    print()
    input("审核完成后按 Enter 继续...")


def step4_merge_and_retrain(reviewed_labels_dir):
    """合并原始标注 + 审核后的伪标签，重新训练"""
    import shutil
    from ultralytics import YOLO

    # 创建新的训练集
    new_train_dir = "Data/Detection_dataset_v2"
    os.makedirs(f"{new_train_dir}/images/train", exist_ok=True)
    os.makedirs(f"{new_train_dir}/labels/train", exist_ok=True)

    # 复制原始训练集
    print("复制原始训练集...")
    shutil.copytree(
        "Data/Detection_dataset/images/train",
        f"{new_train_dir}/images/train",
        dirs_exist_ok=True
    )
    shutil.copytree(
        "Data/Detection_dataset/labels/train",
        f"{new_train_dir}/labels/train",
        dirs_exist_ok=True
    )

    # 添加审核后的伪标签
    print("添加审核后的伪标签...")
    reviewed_imgs = os.listdir(reviewed_labels_dir.replace("/labels/", "/images/"))
    for img_file in reviewed_imgs:
        shutil.copy(
            f"{reviewed_labels_dir.replace('/labels/', '/images/')}/{img_file}",
            f"{new_train_dir}/images/train/{img_file}"
        )
        lbl_file = img_file.replace(".jpg", ".txt")
        shutil.copy(
            f"{reviewed_labels_dir}/{lbl_file}",
            f"{new_train_dir}/labels/train/{lbl_file}"
        )

    print(f"✓ 新训练集创建完成: {new_train_dir}")
    print(f"  原始: 1536 张")
    print(f"  新增: {len(reviewed_imgs)} 张")
    print(f"  总计: {1536 + len(reviewed_imgs)} 张")

    # 更新 dataset.yaml
    yaml_content = f"""
path: {os.path.abspath(new_train_dir)}
train: images/train
val: ../Detection_dataset/images/val

nc: 1
names: ['goat']
"""
    with open(f"{new_train_dir}/dataset.yaml", "w") as f:
        f.write(yaml_content)

    # 重新训练
    print("\n开始重新训练...")
    from Scripts.modules.cbam import register_cbam
    from Scripts.modules.wise_iou import patch_wise_iou

    register_cbam()
    patch_wise_iou()

    model = YOLO("Models/yolo11s_cbam_p2.yaml")
    model.load("yolo11s.pt")

    model.train(
        data=f"{new_train_dir}/dataset.yaml",
        epochs=200,
        imgsz=1280,
        batch=4,
        patience=40,
        project="Detection_experiments",
        name="semi_supervised_v1",
        device=0,
    )

    print("✓ 重新训练完成")


def main():
    print("="*70)
    print("【Phase 1: 半监督学习 + 主动学习】")
    print("="*70)
    print()

    # Step 1: 生成伪标签
    print("Step 1: 生成伪标签...")
    pseudo_label_dir = step1_generate_pseudo_labels()

    # Step 2: 主动学习选择
    print("\nStep 2: 主动学习选择高价值样本...")
    selected_file = step2_active_learning_selection(pseudo_label_dir, top_k=500)

    # Step 3: 人工审核
    print("\nStep 3: 人工审核...")
    step3_human_review_tool(selected_file)

    # Step 4: 合并 + 重新训练
    print("\nStep 4: 合并数据集并重新训练...")
    reviewed_labels_dir = input("输入审核后的标签目录: ")
    step4_merge_and_retrain(reviewed_labels_dir)

    print("\n" + "="*70)
    print("Phase 1 完成！")
    print("预期提升: +5~10%")
    print("="*70)


if __name__ == "__main__":
    main()