#!/usr/bin/env python
"""在弱监督框标签上训练 YOLO 检测模型。

    uv run python detect.py --config configs/detect/film.yaml
    uv run python detect.py --config configs/detect/film.yaml --eval-only \
        --checkpoint runs/detect/film/weights/best.pt

数据是在 Label Studio 里修正过的标签，导出并划分 train/val 之后放在 data/ 下的数据集。
框来自 Grad-CAM / SHAP 的弱监督标注再经人工修正。

**预测框约束为互不重叠**（`train.max_overlap`，默认 0）。真值里一处缺陷只有一个框、
零重叠，模型却会输出重叠框，所以验证阶段补一层抑制——早停与 best.pt 的选择都据此进行，
指标才和最终推理一致。设为负数可关掉。
"""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path

import yaml

from engine.detection import filter_overlapping_outputs
from utils import load_config, seed_everything, setup_logging
from utils.device import DeviceNotAvailableError, describe_device, resolve_device


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="薄膜缺陷检测（YOLO）训练/评估")
    parser.add_argument("--config", required=True, help="yaml 配置路径")
    parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE", help="覆盖配置项")
    parser.add_argument("--work-dir", default=None)
    parser.add_argument("--device", default=None, help="auto / cuda / mps / cpu")
    parser.add_argument("--resume", action="store_true", help="从 work_dir 里的 last.pt 续训")
    parser.add_argument("--eval-only", action="store_true", help="只在验证集上评估")
    parser.add_argument("--checkpoint", default=None, help="评估/微调用的权重 .pt")
    return parser.parse_args()


def install_overlap_suppression(max_overlap: float) -> None:
    """给验证器的后处理包一层"框不重叠"约束。

    只能用打补丁的方式：ultralytics 没有对应的开关——端到端模型（yolo26）的 NMS 早退，
    传入的 iou / agnostic_nms 都不参与 IoU 抑制（见 engine/detection.filter_overlapping_outputs
    的说明）。补在 postprocess 之后，训练期验证、早停、best.pt 的选择就都遵守这条约束。
    """
    from ultralytics.models.yolo.detect import DetectionValidator

    original = DetectionValidator.postprocess

    def postprocess(self, preds):
        return filter_overlapping_outputs(original(self, preds), max_iou=max_overlap)

    DetectionValidator.postprocess = postprocess


def to_ultralytics_device(device) -> str | int:
    """torch.device -> ultralytics 的 device 记法（cuda 用序号，其余用字符串）。"""
    if device.type == "cuda":
        return device.index or 0
    return device.type


def main() -> int:
    args = parse_args()
    cfg = load_config(args.config, overrides=args.set)
    if args.work_dir:
        cfg["work_dir"] = args.work_dir
    work_dir = Path(cfg.get("work_dir", "runs/detect"))
    work_dir.mkdir(parents=True, exist_ok=True)

    logger = setup_logging(log_file=work_dir / "detect.log")
    seed = int(cfg.get("seed", 42))
    seed_everything(seed)

    try:
        device = resolve_device(args.device or cfg.get("device", "auto"))
    except DeviceNotAvailableError as exc:
        logger.error("%s", exc)
        return 1
    logger.info("设备: %s", describe_device(device))

    data_yaml = Path(cfg["data"]["yaml"])
    if not data_yaml.exists():
        logger.error(
            "找不到 %s；先把 Label Studio 里修好的标签导出、划分 train/val 到 data/ 下", data_yaml
        )
        return 1
    meta = yaml.safe_load(data_yaml.read_text(encoding="utf-8"))
    logger.info("数据集: %s（%d 类）", data_yaml, meta["nc"])
    if meta.get("train") == meta.get("val"):
        # generate_boxes.py 的产出就不划分，两者指向同一目录
        logger.warning(
            "data.yaml 的 train 与 val 相同（%s）：验证集就是训练集，mAP 会虚高，"
            "早停与 best.pt 的选择也因此失真。正式训练前请先划分数据集",
            meta.get("train"),
        )

    train_cfg = dict(cfg.get("train", {}))
    swanlab_cfg = dict(train_cfg.pop("swanlab", {}))
    # 不是 ultralytics 的参数，得先摘出来，否则它会拒绝未知键
    max_overlap = float(train_cfg.pop("max_overlap", 0.0))

    from ultralytics import YOLO
    from ultralytics import settings as ul_settings

    if max_overlap >= 0:
        install_overlap_suppression(max_overlap)
        logger.info("验证时约束预测框互不重叠（最大 IoU %.2f）", max_overlap)
    else:
        logger.warning("未约束预测框重叠：指标里会计入真值中不可能出现的重叠框")

    # ultralytics 默认把数据集根指到 ~/datasets 之类的全局目录；我们的 data.yaml 里写的是
    # 绝对路径，但它仍会拿 datasets_dir 去拼相对路径，统一指到项目根省得踩坑。
    ul_settings.update({"datasets_dir": str(Path.cwd())})

    weights = args.checkpoint or cfg["model"]["weights"]
    logger.info("模型: %s", weights)
    model = YOLO(weights)

    if swanlab_cfg.pop("enabled", False):
        try:
            from swanlab.integration.ultralytics import add_swanlab_callback

            os.environ.setdefault("SWANLAB_PROJ_NAME", str(swanlab_cfg.get("project", "film-detect")))
            import swanlab

            swanlab.init(
                project=swanlab_cfg.get("project", "film-detect"),
                name=swanlab_cfg.get("name") or work_dir.name,
                mode=swanlab_cfg.get("mode", "local"),
                config=cfg,
                log_dir=str(work_dir / "swanlog"),
                tags=swanlab_cfg.get("tags") or [],
            )
            add_swanlab_callback(model)
            logger.info("SwanLab 已接入 YOLO 回调")
        except Exception as exc:  # 可视化失败不该挡住训练
            logger.warning("SwanLab 接入失败，训练继续: %s", exc)

    # ultralytics 自己管 project/name 目录。project 必须给绝对路径：相对路径会被它再拼到
    # 全局 settings 的 runs_dir 上，产出就跑到 runs/detect/runs/detect/<name> 去了。
    common = {
        "data": str(data_yaml.resolve()),
        "device": to_ultralytics_device(device),
        "project": str(work_dir.resolve().parent),
        "name": work_dir.name,
        "exist_ok": True,
        "seed": seed,
    }

    if args.eval_only:
        metrics = model.val(split="val", plots=True, **common)
    else:
        model.train(resume=args.resume, **common, **train_cfg)
        # train 结束后 ultralytics 已在最优权重上跑过一次 val，这里显式再跑一次拿回指标对象
        metrics = model.val(split="val", plots=True, **common)

    summary = {
        "map50_95": float(metrics.box.map),
        "map50": float(metrics.box.map50),
        "precision": float(metrics.box.mp),
        "recall": float(metrics.box.mr),
        "per_class_map50": {
            meta["names"][int(c)]: float(v)
            for c, v in zip(metrics.box.ap_class_index, metrics.box.ap50, strict=False)
        },
    }
    logger.info(
        "mAP50-95 %.4f | mAP50 %.4f | P %.4f | R %.4f",
        summary["map50_95"], summary["map50"], summary["precision"], summary["recall"],
    )
    (work_dir / "metrics.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    logger.info(
        "曲线与混淆矩阵见 %s（results.png / *_curve.png / confusion_matrix*.png）",
        work_dir.resolve(),
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())