#!/usr/bin/env python
"""用训练好的检测模型在整个数据集上推理，把预测框与真值框画在一起供人工核对。

    uv run python scripts/visualize_detections.py \\
        --weights runs/detect/film/weights/best.pt \\
        --data data/film_yolo/data.yaml \\
        --out runs/detect/film_vis

图上三种颜色：**绿=真值**，**蓝=命中的预测**，**红=没对上真值的预测**（误检）。
只有绿框没有蓝框，就是漏检。匹配规则与 mAP 一致（按置信度贪心、IoU 阈值、类别需一致），
所以图上看到的和指标算的是同一回事。

产出按"结果类型"分目录，直接对着看失败模式：

```
runs/<name>/
├── train/  val/                 # 分开放：训练集的图模型见过，指标不可信
│   ├── ok/<图名>.png            # 全部命中，没有多余预测
│   ├── missed/<图名>.png        # 有漏检
│   ├── extra/<图名>.png         # 有误检（含背景图上凭空框出的）
│   └── none/<图名>.png          # 既无真值也无预测——负样本判对了
├── index.csv                    # 每张图一行：真值数/预测数/命中数/最高置信度
└── summary.json                 # 逐类 TP/FP/FN 与各结果类型的张数
```

`--only errors` 只输出有漏检或误检的图，扫起来快得多。

**预测框默认约束为互不重叠**（`--max-overlap 0`）。人工标注的真值里一处缺陷只有一个框、
零重叠，而 NMS 只在同类内按固定阈值抑制，留得下跨类重复和 IoU 低于阈值的重叠框。
`--max-overlap -1` 可关掉这层约束，看模型的原始输出。
"""

from __future__ import annotations

import argparse
import csv
import json
import shutil
import sys
from collections import Counter, defaultdict
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import yaml  # noqa: E402

from datasets.yolo import IMAGE_SUFFIXES  # noqa: E402
from engine.detection import (  # noqa: E402
    Box,
    draw_boxes,
    match_boxes,
    parse_label,
    suppress_overlaps,
)
from utils import setup_logging  # noqa: E402
from utils.device import DeviceNotAvailableError, describe_device, resolve_device  # noqa: E402

logger = None

OUTCOMES = ("ok", "missed", "extra", "none")


def outcome(n_gt: int, n_pred: int, match) -> str:
    """一张图的结果类型。missed 优先于 extra：漏检比误检更该先看。"""
    if n_gt == 0 and n_pred == 0:
        return "none"
    if match.unmatched_ground_truth:
        return "missed"
    if match.unmatched_predictions:
        return "extra"
    return "ok"


def load_dataset(data_yaml: Path) -> tuple[list[str], dict[str, list[Path]]]:
    """从 data.yaml 读类别与各 split 的图片列表。"""
    meta = yaml.safe_load(data_yaml.read_text(encoding="utf-8"))
    names = meta["names"]
    classes = [names[i] for i in sorted(names)] if isinstance(names, dict) else list(names)
    root = Path(meta.get("path", data_yaml.parent))
    if not root.is_absolute():
        root = (data_yaml.parent / root).resolve()

    splits = {}
    for split in ("train", "val"):
        entry = meta.get(split)
        if not entry:
            continue
        directory = root / entry
        if not directory.is_dir():
            raise FileNotFoundError(f"{data_yaml} 里的 {split} 指向 {directory}，但它不存在")
        images = sorted(p for p in directory.iterdir() if p.suffix.lower() in IMAGE_SUFFIXES)
        # train 与 val 指向同一目录时（未划分的数据集）只跑一次
        if any(directory == root / meta.get(other, "") for other in splits):
            continue
        splits[split] = images
    if not splits:
        raise FileNotFoundError(f"{data_yaml} 里没有可用的 train / val")
    return classes, splits


def label_path_for(image: Path) -> Path:
    """YOLO 约定：把路径里的 images/ 换成 labels/，后缀换成 .txt。"""
    parts = list(image.parts)
    for index in range(len(parts) - 1, -1, -1):
        if parts[index] == "images":
            parts[index] = "labels"
            break
    return Path(*parts).with_suffix(".txt")


def main() -> int:
    global logger
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--weights", required=True, help="训练好的 .pt")
    parser.add_argument("--data", required=True, help="数据集的 data.yaml")
    parser.add_argument("--out", required=True, help="产出目录，建议放在 runs/ 下")
    parser.add_argument("--device", default="auto", help="auto / cuda / mps / cpu")
    parser.add_argument("--imgsz", type=int, default=640)
    parser.add_argument("--conf", type=float, default=0.25, help="置信度阈值")
    parser.add_argument("--iou-threshold", type=float, default=0.5, help="判定命中的 IoU")
    parser.add_argument(
        "--max-overlap", type=float, default=0.0,
        help="预测框之间允许的最大 IoU，超出的丢低分框。0=完全不允许重叠（默认，与真值一致），"
             "-1=不约束，看模型原始输出",
    )
    parser.add_argument("--nms-iou", type=float, default=0.7, help="NMS 的 IoU 阈值")
    parser.add_argument(
        "--no-agnostic-nms", action="store_true",
        help="关掉类别无关 NMS。默认开启，用来压掉同一处缺陷被两个类别各框一次的重复",
    )
    parser.add_argument("--batch", type=int, default=16)
    parser.add_argument(
        "--only", choices=("all", "errors"), default="all",
        help="errors=只输出有漏检或误检的图，统计仍按全量计",
    )
    parser.add_argument("--limit", type=int, default=0, help="每个 split 只跑前 N 张（试参数用）")
    parser.add_argument("--overwrite", action="store_true", help="产出目录已存在时先清空")
    args = parser.parse_args()

    out_dir = Path(args.out)
    logger = setup_logging()

    try:
        device = resolve_device(args.device)
        classes, splits = load_dataset(Path(args.data))
    except (DeviceNotAvailableError, FileNotFoundError) as exc:
        logger.error("%s", exc)
        return 1

    if args.limit:
        splits = {name: images[: args.limit] for name, images in splits.items()}
    total = sum(len(images) for images in splits.values())
    logger.info("设备: %s | 权重: %s", describe_device(device), args.weights)
    logger.info(
        "数据集: %s（%d 类，%s）", args.data, len(classes),
        "，".join(f"{name} {len(images)} 张" for name, images in splits.items()),
    )
    logger.info(
        "conf=%.2f  命中 IoU=%.2f  imgsz=%d  NMS IoU=%.2f  类别无关 NMS=%s  框间最大重叠=%s",
        args.conf, args.iou_threshold, args.imgsz, args.nms_iou,
        not args.no_agnostic_nms,
        "不约束" if args.max_overlap < 0 else f"{args.max_overlap:.2f}",
    )

    if out_dir.exists() and any(out_dir.iterdir()):
        if not args.overwrite:
            logger.error("%s 已存在且非空；确认要覆盖请加 --overwrite", out_dir)
            return 1
        shutil.rmtree(out_dir)

    from PIL import Image
    from ultralytics import YOLO

    model = YOLO(args.weights)
    device_arg = device.index or 0 if device.type == "cuda" else device.type

    outcomes: dict[str, Counter[str]] = {name: Counter() for name in splits}
    stats: dict[str, dict[str, Counter[str]]] = {
        name: {"tp": Counter(), "fp": Counter(), "fn": Counter()} for name in splits
    }
    scores: dict[str, list[float]] = defaultdict(list)
    rows = []
    processed = 0
    suppressed = 0

    def predict_in_batches(images):
        """自己按批切分：ultralytics 的 predict 不理会 batch 参数，
        一次性把整个列表当成一个 batch，上万张图必然 OOM。"""
        for start in range(0, len(images), args.batch):
            chunk = images[start:start + args.batch]
            results = model.predict(
                source=[str(p) for p in chunk], imgsz=args.imgsz, conf=args.conf,
                device=device_arg, verbose=False, iou=args.nms_iou,
                agnostic_nms=not args.no_agnostic_nms,
            )
            yield from zip(chunk, results, strict=True)

    for split, images in splits.items():
        for image_path, result in predict_in_batches(images):
            predictions = [
                Box(int(cls), *(float(v) for v in xyxy), float(score))
                for xyxy, cls, score in zip(
                    result.boxes.xyxyn.tolist(), result.boxes.cls.tolist(),
                    result.boxes.conf.tolist(), strict=True,
                )
            ]
            if args.max_overlap >= 0:
                kept = suppress_overlaps(predictions, max_iou=args.max_overlap)
                suppressed += len(predictions) - len(kept)
                predictions = [predictions[i] for i in kept]

            label_file = label_path_for(image_path)
            truth = parse_label(label_file.read_text(encoding="utf-8")) if label_file.exists() else []
            match = match_boxes(predictions, truth, args.iou_threshold)

            kind = outcome(len(truth), len(predictions), match)
            outcomes[split][kind] += 1
            for _, truth_index in match.pairs:
                stats[split]["tp"][classes[truth[truth_index].class_id]] += 1
            for index in match.unmatched_predictions:
                stats[split]["fp"][classes[predictions[index].class_id]] += 1
            for index in match.unmatched_ground_truth:
                stats[split]["fn"][classes[truth[index].class_id]] += 1
            for box in predictions:
                scores[classes[box.class_id]].append(box.score)

            rows.append({
                "split": split, "name": image_path.name, "outcome": kind,
                "n_gt": len(truth), "n_pred": len(predictions), "n_matched": len(match.pairs),
                "max_conf": round(max((b.score for b in predictions), default=0.0), 4),
                "gt_classes": "|".join(sorted({classes[b.class_id] for b in truth})),
                "pred_classes": "|".join(sorted({classes[b.class_id] for b in predictions})),
            })

            if args.only == "all" or kind in ("missed", "extra"):
                destination = out_dir / split / kind / f"{image_path.stem}.png"
                destination.parent.mkdir(parents=True, exist_ok=True)
                with Image.open(image_path) as raw:
                    canvas = raw.convert("RGB")
                draw_boxes(canvas, truth, predictions, classes, match)
                canvas.save(destination)

            processed += 1
            if processed % 1000 == 0:
                logger.info("处理 [%d/%d]", processed, total)

    with (out_dir / "index.csv").open("w", encoding="utf-8-sig", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)

    summary = {
        "weights": str(Path(args.weights).resolve()),
        "data": str(Path(args.data).resolve()),
        "conf": args.conf,
        "iou_threshold": args.iou_threshold,
        "imgsz": args.imgsz,
        "nms_iou": args.nms_iou,
        "agnostic_nms": not args.no_agnostic_nms,
        "max_overlap": args.max_overlap,
        "suppressed_overlapping_boxes": suppressed,
        "outcomes": {split: dict(counter) for split, counter in outcomes.items()},
        "per_class": {
            split: {
                name: {
                    "tp": counters["tp"][name], "fp": counters["fp"][name],
                    "fn": counters["fn"][name],
                }
                for name in classes
            }
            for split, counters in stats.items()
        },
    }
    (out_dir / "summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8"
    )

    if suppressed:
        logger.info("因与更高分的框重叠而丢弃: %d 个", suppressed)
    for split, counter in outcomes.items():
        shown = sum(counter[k] for k in ("missed", "extra"))
        logger.info(
            "%s：全部命中 %d，漏检 %d，误检 %d，负样本判对 %d（有问题的 %d 张，占 %.1f%%）",
            split, counter["ok"], counter["missed"], counter["extra"], counter["none"],
            shown, 100 * shown / max(sum(counter.values()), 1),
        )
        logger.info("%-12s %8s %8s %8s %8s", "类别", "TP", "FP", "FN", "召回")
        for name in classes:
            tp, fp, fn = (stats[split][k][name] for k in ("tp", "fp", "fn"))
            if tp or fp or fn:
                logger.info(
                    "%-12s %8d %8d %8d %7.1f%%",
                    name, tp, fp, fn, 100 * tp / max(tp + fn, 1),
                )
    logger.info("产出目录: %s（index.csv 可按 outcome / max_conf 排序筛选）", out_dir)
    return 0


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