#!/usr/bin/env python
"""用训练好的分类模型 + Grad-CAM 为每张图生成框标签。

    uv run python scripts/gradcam_boxes.py --checkpoint runs/classify/best.pth

分类模型只学过"这是什么缺陷"。Grad-CAM 把分类分数对卷积特征的梯度当作通道重要性，
加权求和得到空间热力图——这是分类模型隐含的定位线索。阈值化 + 取连通域外接矩形，
就得到弱监督的框标签。

产出：

```
runs/gradcam_boxes/
├── labels/<类别>/<图名>.txt   # YOLO 格式：class_id cx cy w h（归一化）
├── viz/<类别>/<图名>.png      # 原图 | 热力图 | 叠加+框，供人工核对匹配程度
├── data.yaml                  # 类别表，可直接喂给检测训练
└── summary.json
```

**框的类别取自图片所在目录**（人工归类的结果），不是模型预测——即使模型分错，
我们要的也是该图真实类别的热力区域。

分辨率是主要限制：resnet18 在 128×128 输入下 layer4 只有 4×4、layer3 有 8×8。
`--stage` 可选层（4~7 对应 layer1~4），层越浅越精细但类别判别性越弱。
"""

from __future__ import annotations

import argparse
import json
import sys
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


def save_visualization(path: Path, image_path: Path, cam: np.ndarray, boxes, title: str) -> None:
    """三联图：原图 | 热力图 | 叠加+框。人工一眼就能看出热力区域与框合不合。"""
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib.patches import Rectangle

    with Image.open(image_path) as raw:
        image = np.asarray(raw.convert("RGB"))
    height, width = image.shape[:2]
    heat = np.asarray(
        Image.fromarray((cam * 255).astype(np.uint8)).resize((width, height), Image.BILINEAR)
    ) / 255.0

    fig, axes = plt.subplots(1, 3, figsize=(11, 4))
    axes[0].imshow(image)
    axes[0].set_title("original")
    axes[1].imshow(heat, cmap="jet", vmin=0, vmax=1)
    axes[1].set_title(f"grad-cam {cam.shape[0]}x{cam.shape[1]}")
    axes[2].imshow(image)
    axes[2].imshow(heat, cmap="jet", alpha=0.45, vmin=0, vmax=1)
    for x1, y1, x2, y2 in boxes:
        axes[2].add_patch(
            Rectangle(
                (x1 * width, y1 * height), (x2 - x1) * width, (y2 - y1) * height,
                fill=False, edgecolor="lime", linewidth=2,
            )
        )
    axes[2].set_title(f"overlay + {len(boxes)} box")
    for ax in axes:
        ax.axis("off")
    fig.suptitle(title, fontsize=10)
    fig.tight_layout()
    path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(path, dpi=110)
    plt.close(fig)


def main() -> int:
    global logger
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--checkpoint", default="runs/classify/best.pth", help="分类模型权重")
    parser.add_argument("--config", default="configs/classify/film.yaml", help="训练时的配置")
    parser.add_argument("--out", default="runs/gradcam_boxes", help="输出目录")
    parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE")
    parser.add_argument("--device", default=None)
    parser.add_argument(
        "--stage", type=int, default=6,
        help="取骨架 features 的第几层做 CAM（4~7 = layer1~4）；默认 6 即 layer3",
    )
    parser.add_argument("--threshold", type=float, default=0.4, help="热力图阈值（相对最大值）")
    parser.add_argument("--padding", type=float, default=0.1, help="框四周外扩比例")
    parser.add_argument(
        "--max-boxes", type=int, default=1,
        help="每张图最多保留几个框（按连通域内热力总和排序取前 N 个）。"
             "默认 1：一张裁片通常只对应一处缺陷，多出来的框多半是噪声",
    )
    parser.add_argument(
        "--min-area-ratio", type=float, default=0.02, help="连通域面积下限（占整图比例）"
    )
    parser.add_argument(
        "--peak-ratio", type=float, default=0.75,
        help="次要连通域的峰值至少要达到全局最大值的这个比例才保留，用于滤掉弱响应假框",
    )
    parser.add_argument(
        "--viz-per-class", type=int, default=20, help="每个类别存多少张可视化，0 = 不存"
    )
    parser.add_argument("--viz-all", action="store_true", help="每张图都存可视化（很慢很占盘）")
    parser.add_argument("--limit", type=int, default=0, help="只处理前 N 张，用于快速试参数")
    args = parser.parse_args()

    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    logger = setup_logging(out_dir / "gradcam.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 [],
    )
    if args.limit:
        dataset.samples = dataset.samples[: args.limit]
    classes = dataset.classes

    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():      # Grad-CAM 需要梯度，但不更新权重
        param.requires_grad_(True)

    logger.info("设备: %s | 权重: %s", describe_device(device), args.checkpoint)
    logger.info("数据集 %s：%d 张，%d 类", data_cfg["root"], len(dataset), len(classes))
    logger.info(
        "CAM 目标层 features[%d] | 阈值 %.2f | 外扩 %.2f | 每图最多 %d 框",
        args.stage, args.threshold, args.padding, args.max_boxes,
    )

    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,
    )

    labels_dir = out_dir / "labels"
    viz_dir = out_dir / "viz"
    viz_done: Counter[str] = Counter()
    box_counts: Counter[int] = Counter()
    per_class_boxes: Counter[str] = Counter()
    correct = 0
    cam_shape = None

    for step, batch in enumerate(loader):
        images = batch["images"].to(device)
        targets = batch["labels"].to(device)
        cams = grad_cam(model.backbone, model.head, images, targets, stage=args.stage)
        with torch.no_grad():
            correct += int((model(images).argmax(1) == targets).sum())
        cam_shape = tuple(cams.shape[-2:])

        for cam, target, meta in zip(cams.cpu().numpy(), targets.cpu().tolist(), batch["meta"],
                                     strict=True):
            name = classes[target]
            boxes = cam_to_boxes(
                cam,
                threshold=args.threshold,
                min_area_ratio=args.min_area_ratio,
                max_boxes=args.max_boxes,
                padding=args.padding,
                peak_ratio=args.peak_ratio,
            )
            box_counts[len(boxes)] += 1
            per_class_boxes[name] += len(boxes)

            source = Path(meta["path"])
            label_path = labels_dir / name / f"{source.stem}.txt"
            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 args.viz_all or viz_done[name] < args.viz_per_class:
                viz_done[name] += 1
                save_visualization(
                    viz_dir / name / f"{source.stem}.png", source, cam, boxes,
                    f"{name}  |  {len(boxes)} box  |  thr={args.threshold}",
                )

        if step % 20 == 0:
            logger.info("处理 [%d/%d]", step, len(loader))

    total = sum(box_counts.values())
    write_data_yaml(out_dir / "data.yaml", classes, root=out_dir)
    summary = {
        "checkpoint": args.checkpoint,
        "dataset": str(data_cfg["root"]),
        "num_images": total,
        "cam_shape": list(cam_shape) if cam_shape else None,
        "stage": args.stage,
        "threshold": args.threshold,
        "padding": args.padding,
        "max_boxes": args.max_boxes,
        "min_area_ratio": args.min_area_ratio,
        "peak_ratio": args.peak_ratio,
        "classification_accuracy_on_all": correct / total if total else 0.0,
        "boxes_per_image": {str(k): v for k, v in sorted(box_counts.items())},
        "boxes_per_class": dict(sorted(per_class_boxes.items())),
    }
    (out_dir / "summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8"
    )

    logger.info("CAM 分辨率 %s | 全量分类准确率 %.4f", cam_shape, summary["classification_accuracy_on_all"])
    for count, number in sorted(box_counts.items()):
        logger.info("  生成 %d 个框的图片: %d 张 (%.1f%%)", count, number, 100 * number / total)
    logger.info("标签写入 %s，可视化 %d 张写入 %s", labels_dir, sum(viz_done.values()), viz_dir)
    return 0


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