"""Grad-CAM 与由热力图导出检测框。

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

分辨率是主要限制：resnet18 在 128×128 输入下 layer4 只有 4×4、layer3 有 8×8。
层越浅空间越精细，但类别判别性越弱，需要按数据实测取舍。
"""

from __future__ import annotations

import numpy as np
import torch
from torch import nn
from torch.nn import functional as F


def grad_cam(
    backbone: nn.Module,
    head: nn.Module,
    images: torch.Tensor,
    targets: torch.Tensor,
    stage: int | None = None,
) -> torch.Tensor:
    """对一个 batch 求 Grad-CAM，返回 (B, h, w) 的热力图（每张已归一化到 [0, 1]）。

    Args:
        stage: 取 `backbone.features` 的第几层输出做 CAM；None 表示用骨架最终输出
            （即最后一个卷积块）。层越浅分辨率越高、类别判别性越弱。

    对每个样本只回传它自己那一类的分数，所以整批可以一次 backward——
    样本之间的梯度互不干扰。
    """
    layers = getattr(backbone, "features", None)
    if stage is not None and not isinstance(layers, nn.Sequential):
        raise ValueError("该骨架不是 nn.Sequential，无法按 stage 取中间层")

    with torch.enable_grad():
        if stage is None:
            activations = backbone(images)
            features = activations
        else:
            features = images
            for index, layer in enumerate(layers):
                features = layer(features)
                if index == stage:
                    break
            else:
                raise ValueError(f"stage {stage} 超出骨架层数 {len(layers)}")
            activations = features
            for layer in list(layers)[stage + 1 :]:
                features = layer(features)

        activations.retain_grad()
        logits = head(features)
        # 只取每个样本对应目标类别的分数；求和后一次 backward 即可
        score = logits.gather(1, targets.view(-1, 1)).sum()
        backbone.zero_grad(set_to_none=True)
        head.zero_grad(set_to_none=True)
        score.backward()

    gradients = activations.grad
    if gradients is None:
        raise RuntimeError("没有拿到梯度，检查目标层是否参与了前向计算")

    weights = gradients.mean(dim=(2, 3), keepdim=True)   # 各通道的重要性
    cam = F.relu((weights * activations).sum(dim=1))     # (B, h, w)

    flat = cam.flatten(1)
    lo = flat.min(dim=1, keepdim=True).values
    hi = flat.max(dim=1, keepdim=True).values
    cam = ((flat - lo) / (hi - lo).clamp(min=1e-8)).view_as(cam)
    return cam.detach()


def connected_components(mask: np.ndarray) -> list[np.ndarray]:
    """4-邻接连通域标记，返回每个连通域的布尔掩码。

    自己写而不引 scipy/opencv：项目只依赖 numpy/torch/PIL，
    为这一处再加一个重依赖不划算。
    """
    height, width = mask.shape
    seen = np.zeros_like(mask, dtype=bool)
    components = []
    for start_y in range(height):
        for start_x in range(width):
            if not mask[start_y, start_x] or seen[start_y, start_x]:
                continue
            stack = [(start_y, start_x)]
            seen[start_y, start_x] = True
            pixels = []
            while stack:
                y, x = stack.pop()
                pixels.append((y, x))
                for ny, nx in ((y - 1, x), (y + 1, x), (y, x - 1), (y, x + 1)):
                    if 0 <= ny < height and 0 <= nx < width and mask[ny, nx] and not seen[ny, nx]:
                        seen[ny, nx] = True
                        stack.append((ny, nx))
            component = np.zeros_like(mask, dtype=bool)
            ys, xs = zip(*pixels, strict=True)
            component[list(ys), list(xs)] = True
            components.append(component)
    return components


def cam_to_boxes(
    cam: np.ndarray,
    threshold: float = 0.4,
    min_area_ratio: float = 0.01,
    max_boxes: int = 1,
    padding: float = 0.0,
    peak_ratio: float = 0.75,
) -> list[tuple[float, float, float, float]]:
    """把一张热力图变成若干归一化 xyxy 框。

    Args:
        threshold: 相对最大值的阈值；CAM 已归一化到 [0,1]，所以直接比较。
        min_area_ratio: 连通域面积小于整图这个比例就丢掉，滤掉零星的高响应点。
        max_boxes: 最多保留几个框（按连通域内热力总和排序）。默认 1——一张裁片通常只对应一处缺陷。
        padding: 框的四周外扩比例——CAM 倾向只覆盖最具判别性的部位（比如缺陷中心），
            往往比真实缺陷范围小，适当外扩更贴合。
        peak_ratio: 连通域的峰值至少要达到全局最大值的这个比例才保留。
            CAM 归一化后最强的区域峰值恒为 1，次强区域若只有 0.45 多半是噪声——
            实测不加这条会在没有缺陷的角落刷出假框。

    返回的坐标已归一化到 [0, 1]，与图像实际尺寸无关，可直接写 YOLO 标签。
    """
    height, width = cam.shape
    mask = cam >= threshold
    if not mask.any():
        return []

    boxes = []
    for component in connected_components(mask):
        if component.sum() < max(1, min_area_ratio * height * width):
            continue
        if cam[component].max() < peak_ratio:   # 弱响应区多半是噪声，不是第二个缺陷
            continue
        ys, xs = np.where(component)
        # +1 让框覆盖住整个格子，而不是格子中心之间的距离
        x1, x2 = xs.min() / width, (xs.max() + 1) / width
        y1, y2 = ys.min() / height, (ys.max() + 1) / height
        if padding:
            dx, dy = (x2 - x1) * padding, (y2 - y1) * padding
            x1, x2 = max(0.0, x1 - dx), min(1.0, x2 + dx)
            y1, y2 = max(0.0, y1 - dy), min(1.0, y2 + dy)
        boxes.append((float(cam[component].sum()), (x1, y1, x2, y2)))

    boxes.sort(key=lambda item: -item[0])
    return [box for _, box in boxes[:max_boxes]]


def to_yolo_line(box: tuple[float, float, float, float], class_id: int) -> str:
    """归一化 xyxy -> YOLO 的 `class_id cx cy w h`。"""
    x1, y1, x2, y2 = box
    return (
        f"{class_id} {(x1 + x2) / 2:.6f} {(y1 + y2) / 2:.6f} {x2 - x1:.6f} {y2 - y1:.6f}"
    )