#!/usr/bin/env python
"""用训练好的分类模型为每张图生成热力图与 YOLO 框标签。

    # Grad-CAM（快，全量 11k 张约 2 分钟）
    uv run python scripts/generate_boxes.py --out runs/boxes_gradcam --method gradcam

    # SHAP Partition（慢但更精细，全量约 1.7 小时）
    uv run python scripts/generate_boxes.py --out runs/boxes_shap --method shap

两种归因方法都能定位，取舍不同：

| | gradcam | shap |
| --- | --- | --- |
| 分辨率 | 受限于特征图（layer3 @128² 只有 8×8） | 与输入同分辨率（128×128） |
| 单张耗时 | 一次 backward，约 0.01 秒 | 几百次前向，约 0.55 秒 |
| 小目标 | 几像素的黑点在 8×8 上不可分辨 | 能贴合到几十像素的目标 |

产出是一个**自成一体的 YOLO 数据集**，可直接喂给 train.py：

```
runs/<name>/
├── images/<类别>/<图名>.png   # 硬链接到原图，不占额外磁盘
├── labels/<类别>/<图名>.txt   # YOLO：class_id cx cy w h（归一化）
├── heatmaps/<类别>/<图名>.png # 热力图叠加原图 + 框，供人工核对
├── data.yaml
└── summary.json
```

images/ 与 labels/ 的对应关系符合本仓库的 YOLO 约定（datasets/yolo.py），
所以 `FilmDefectDataset` 能直接读。

默认对**全量**数据生成；`--limit` / `--per-class` 只用于快速试参数。
热力图默认针对模型的 **top1 预测类别**；`--target label` 可改用图片所在目录的标注类别。
"""

from __future__ import annotations

import argparse
import errno
import json
import os
import random
import shutil
import sys
import time
from collections import Counter
from pathlib import Path

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

import numpy as np  # noqa: E402
import torch  # noqa: E402
from PIL import Image  # noqa: E402
from torch.utils.data import DataLoader  # noqa: E402

from datasets import FolderClassificationDataset, classification_collate  # noqa: E402
from datasets.yolo import write_data_yaml  # noqa: E402
from engine import load_checkpoint  # noqa: E402
from engine.gradcam import cam_to_boxes, grad_cam, to_yolo_line  # noqa: E402
from models import build_model  # noqa: E402
from utils import load_config, setup_logging  # noqa: E402
from utils.device import DeviceNotAvailableError, describe_device, resolve_device  # noqa: E402

logger = None

METHODS = ("gradcam", "shap")

#: 两种方法的归因图分辨率差 16 倍，面积阈值不能共用一个默认值：
#: 同一个小黑点在 8×8 上占几格、在 128×128 上只占百来像素，
#: 沿用 0.02（=328 像素）会把它整个滤掉。
DEFAULT_MIN_AREA_RATIO = {"gradcam": 0.02, "shap": 0.002}


class NameCollision(RuntimeError):
    """两张不同的源图在产出目录里撞到了同一个文件名。"""


def link_or_copy(source: Path, destination: Path) -> str:
    """给原图建硬链接；跨文件系统时退回复制。

    硬链接不占额外磁盘，却表现得和真实文件一样——产出目录可以整个拷走。
    """
    if destination.exists():
        return "exists"
    destination.parent.mkdir(parents=True, exist_ok=True)
    try:
        os.link(source, destination)
        return "hard"
    except OSError as exc:
        if exc.errno != errno.EXDEV:
            raise
        shutil.copy2(source, destination)
        return "copy"


def save_heatmap(path: Path, image_path: Path, heat: np.ndarray, boxes) -> None:
    """热力图叠加原图并画框。用 PIL 直接合成而不是 matplotlib：
    全量 11k 张要出图，逐张建 figure 太慢也太占盘。"""
    from matplotlib import colormaps

    with Image.open(image_path) as raw:
        base = raw.convert("RGB")
    width, height = base.size

    resized = np.asarray(
        Image.fromarray((heat * 255).astype(np.uint8)).resize((width, height), Image.BILINEAR)
    ) / 255.0
    colored = (colormaps["jet"](resized)[..., :3] * 255).astype(np.uint8)
    blended = Image.blend(base, Image.fromarray(colored), alpha=0.45)

    from PIL import ImageDraw

    draw = ImageDraw.Draw(blended)
    for x1, y1, x2, y2 in boxes:
        draw.rectangle(
            [x1 * width, y1 * height, x2 * width, y2 * height], outline=(0, 255, 0), width=2
        )
    path.parent.mkdir(parents=True, exist_ok=True)
    blended.save(path)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--out", required=True, help="产出目录，建议放在 runs/ 下")
    parser.add_argument(
        "--method", choices=METHODS, default="gradcam", help="归因方法，默认 gradcam"
    )
    parser.add_argument("--checkpoint", default="runs/classify/best.pth", help="分类模型权重")
    parser.add_argument("--config", default="configs/classify/film.yaml", help="训练时的配置")
    parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE")
    parser.add_argument("--device", default=None)
    parser.add_argument(
        "--target", choices=("predicted", "label"), default="predicted",
        help="热力图针对哪个类别：predicted=模型 top1 预测（默认），label=图片所在目录",
    )

    parser.add_argument("--threshold", type=float, default=0.4, help="热力图阈值（相对最大值）")
    parser.add_argument(
        "--max-boxes", type=int, default=1,
        help="每张图最多保留几个框，按归因总量排序取前 N 个。"
             "默认 1：一张裁片通常只对应一处缺陷",
    )
    parser.add_argument("--padding", type=float, default=0.1, help="框四周外扩比例")
    parser.add_argument(
        "--min-area-ratio", type=float, default=None,
        help=f"连通域面积下限（占整图比例）。缺省随方法而定：{DEFAULT_MIN_AREA_RATIO}",
    )
    parser.add_argument(
        "--peak-ratio", type=float, default=0.75,
        help="次要连通域的峰值需达到全局最大值的这个比例才保留，用于滤掉弱响应假框",
    )

    parser.add_argument(
        "--stage", type=int, default=6,
        help="gradcam 的目标层（4~7 = layer1~4）。默认 6 即 layer3（8×8）；"
             "layer4 只有 4×4 太粗，layer2 响应太散",
    )
    parser.add_argument("--max-evals", type=int, default=500, help="shap 每张图的前向次数")
    parser.add_argument("--smooth-radius", type=int, default=2, help="shap 归因图平滑半径，0=不平滑")

    parser.add_argument("--no-heatmaps", action="store_true", help="不输出热力图，只出标签")
    parser.add_argument("--no-images", action="store_true", help="不建原图硬链接")
    parser.add_argument("--limit", type=int, default=0, help="只处理前 N 张（试参数用）")
    parser.add_argument(
        "--per-class", type=int, default=0, help="每类只处理 N 张（试参数用），0=全量"
    )
    parser.add_argument("--seed", type=int, default=0, help="--per-class 抽样的随机种子")
    return parser.parse_args()


def main() -> int:
    global logger
    args = parse_args()
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    logger = setup_logging(out_dir / "generate_boxes.log")

    cfg = load_config(args.config, args.set)
    try:
        device = resolve_device(args.device or cfg.get("device", "auto"))
    except DeviceNotAvailableError as exc:
        logger.error("%s", exc)
        return 1

    data_cfg = cfg["data"]
    dataset = FolderClassificationDataset(
        root=data_cfg["root"], split="all", classes=data_cfg.get("classes"),
        transforms=data_cfg["val_transforms"], exclude=data_cfg.get("exclude") or [],
    )
    classes = dataset.classes

    if args.per_class:
        grouped: dict[int, list] = {}
        for sample in dataset.samples:
            grouped.setdefault(sample[1], []).append(sample)
        rng = random.Random(args.seed)
        dataset.samples = [
            s for group in grouped.values() for s in rng.sample(group, min(args.per_class, len(group)))
        ]
    if args.limit:
        dataset.samples = dataset.samples[: args.limit]

    model_cfg = dict(cfg["model"])
    model_cfg.setdefault("num_classes", len(classes))
    model_cfg["init_from"] = None          # 权重来自 checkpoint
    model = build_model(model_cfg)
    load_checkpoint(args.checkpoint, model, map_location=device)
    model.to(device).eval()
    for param in model.parameters():       # 归因需要梯度，但不更新权重
        param.requires_grad_(True)

    min_area_ratio = (
        args.min_area_ratio if args.min_area_ratio is not None
        else DEFAULT_MIN_AREA_RATIO[args.method]
    )
    box_kwargs = {
        "threshold": args.threshold, "min_area_ratio": min_area_ratio,
        "max_boxes": args.max_boxes, "padding": args.padding, "peak_ratio": args.peak_ratio,
    }

    logger.info("设备: %s | 权重: %s", describe_device(device), args.checkpoint)
    logger.info("方法: %s | 目标类别: %s", args.method, args.target)
    logger.info("数据集 %s：%d 张，%d 类", data_cfg["root"], len(dataset), len(classes))
    logger.info("框参数: %s", box_kwargs)
    if args.method == "shap":
        logger.warning(
            "shap 单张约 0.55 秒，%d 张预计 %.1f 分钟；只想试参数可用 --per-class / --limit",
            len(dataset), len(dataset) * 0.55 / 60,
        )

    started = time.time()
    box_counts: Counter[int] = Counter()
    link_modes: Counter[str] = Counter()
    seen: dict[Path, Path] = {}
    agree = 0
    heat_shape = None

    if args.method == "gradcam":
        loader = DataLoader(
            dataset,
            batch_size=int(data_cfg.get("loader", {}).get("batch_size", 64)),
            shuffle=False,
            num_workers=int(data_cfg.get("loader", {}).get("num_workers", 4)),
            collate_fn=classification_collate,
        )
        processed = 0
        for step, batch in enumerate(loader):
            images = batch["images"].to(device)
            labels = batch["labels"].to(device)
            with torch.no_grad():
                predicted = model(images).argmax(1)
            targets = predicted if args.target == "predicted" else labels
            agree += int((predicted == labels).sum())
            heats = grad_cam(model.backbone, model.head, images, targets, stage=args.stage)
            heat_shape = tuple(heats.shape[-2:])

            for heat, target, meta in zip(
                heats.cpu().numpy(), targets.cpu().tolist(), batch["meta"], strict=True
            ):
                try:
                    _emit(out_dir, Path(meta["path"]), classes, target, heat, box_kwargs,
                          args, box_counts, link_modes, seen)
                except NameCollision as exc:
                    logger.error("%s", exc)
                    return 1
                processed += 1
            if step % 20 == 0:
                logger.info("处理 [%d/%d]", processed, len(dataset))
    else:
        from engine.shap_explain import shap_attribution, smooth

        for index in range(len(dataset)):
            sample = dataset[index]
            image_chw = sample["image"]
            with torch.no_grad():
                predicted = int(model(image_chw[None].to(device)).argmax(1))
            agree += int(predicted == sample["label"])
            target = predicted if args.target == "predicted" else sample["label"]

            heat = shap_attribution(
                model, image_chw.permute(1, 2, 0).numpy().astype(np.float32), target, device,
                max_evals=args.max_evals,
            )
            heat = smooth(heat, args.smooth_radius)
            heat_shape = heat.shape
            try:
                _emit(out_dir, Path(sample["meta"]["path"]), classes, target, heat, box_kwargs,
                      args, box_counts, link_modes, seen)
            except NameCollision as exc:
                logger.error("%s", exc)
                return 1
            if index % 50 == 0:
                elapsed = time.time() - started
                logger.info(
                    "处理 [%d/%d] 已用 %.0fs，平均 %.2fs/张", index, len(dataset),
                    elapsed, elapsed / max(index, 1),
                )

    total = sum(box_counts.values())
    write_data_yaml(out_dir / "data.yaml", classes, root=out_dir,
                    train="images", val="images")
    summary = {
        "method": args.method,
        "checkpoint": args.checkpoint,
        "dataset": str(data_cfg["root"]),
        "target": args.target,
        "num_images": total,
        "heatmap_shape": list(heat_shape) if heat_shape else None,
        "seconds_per_image": (time.time() - started) / max(total, 1),
        "box_params": box_kwargs,
        "stage": args.stage if args.method == "gradcam" else None,
        "max_evals": args.max_evals if args.method == "shap" else None,
        "prediction_matches_folder_label": agree / total if total else 0.0,
        "boxes_per_image": {str(k): v for k, v in sorted(box_counts.items())},
        "link_modes": dict(link_modes),
    }
    (out_dir / "summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8"
    )

    logger.info(
        "完成 %d 张，热力图 %s，平均 %.3fs/张", total, heat_shape, summary["seconds_per_image"]
    )
    logger.info("模型预测与目录标注一致率: %.4f", summary["prediction_matches_folder_label"])
    for count, number in sorted(box_counts.items()):
        logger.info("  生成 %d 个框: %d 张 (%.1f%%)", count, number, 100 * number / total)
    if link_modes:
        logger.info("原图链接方式: %s", dict(link_modes))
    logger.info("产出目录: %s", out_dir)
    return 0


def _emit(out_dir, source, classes, target, heat, box_kwargs, args, box_counts, link_modes,
          seen) -> None:
    """把一张图的结果落盘：硬链接 + YOLO 标签 + 热力图。

    产出按**预测类别**分目录，而源图来自不同的类别目录，所以两张同名的源图
    （dirA/0.png 与 dirB/0.png）可能落到同一个目标路径上——后写的会覆盖前一个的标签。
    这里显式挡住，而不是让标签悄悄丢失。
    """
    name = classes[target]
    boxes = cam_to_boxes(heat, **box_kwargs)
    box_counts[len(boxes)] += 1

    label_path = out_dir / "labels" / name / f"{source.stem}.txt"
    previous = seen.get(label_path)
    if previous is not None:
        raise NameCollision(
            f"{previous} 与 {source} 在产出目录里都要写 {label_path}，后者会覆盖前者。"
            "源图文件名需要全局唯一——用 scripts/export_clusters.py --prefix-parent "
            "重新导出数据集即可。"
        )
    seen[label_path] = source
    label_path.parent.mkdir(parents=True, exist_ok=True)
    label_path.write_text(
        "\n".join(to_yolo_line(box, target) for box in boxes) + ("\n" if boxes else ""),
        encoding="utf-8",
    )
    if not args.no_images:
        link_modes[link_or_copy(source, out_dir / "images" / name / source.name)] += 1
    if not args.no_heatmaps:
        save_heatmap(out_dir / "heatmaps" / name / f"{source.stem}.png", source, heat, boxes)


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