"""用 SHAP 的 Partition explainer 做定位归因。

与 Grad-CAM 的取舍正好互补：

- **分辨率**：Grad-CAM 受限于卷积特征图（resnet18 @128² 的 layer3 只有 8×8），
  SHAP 直接在像素上做层次划分，输出与输入同分辨率。
- **代价**：Grad-CAM 一次 backward 就够；Partition explainer 要几百次前向，
  单张约 1 秒，全量 11k 张需要数小时。

原理：把图像层次化地切块，用 Owen 值近似 Shapley 值，衡量每块对目标类别分数的贡献。
遮挡用数值 masker（默认 0）而不是教程里的 inpaint——我们的图已做过逐图标准化，
0 恰好是该图均值，语义上就是"抹掉这块的信息"，而且不必依赖 opencv。
"""

from __future__ import annotations

from collections.abc import Callable

import numpy as np
import torch
from torch import nn


def shap_attribution(
    model: nn.Module,
    image: np.ndarray,
    target: int,
    device: torch.device,
    max_evals: int = 500,
    batch_size: int = 64,
    mask_value: float = 0.0,
) -> np.ndarray:
    """对单张图求目标类别的 SHAP 归因，返回与输入同尺寸的 (H, W) 图，已归一化到 [0, 1]。

    Args:
        image: HWC 的 float 数组，取值范围与训练时一致（这里是逐图标准化之后的）。
        target: 要解释的类别下标——用图片所在目录的真实类别，而不是模型预测。

    只保留**正贡献**：负贡献表示"这块像素在把分数往下拉"，对定位没有意义。
    """
    import shap

    def predict(batch: np.ndarray) -> np.ndarray:
        tensor = torch.from_numpy(batch).permute(0, 3, 1, 2).float().to(device)
        with torch.no_grad():
            # 只输出目标类别这一列，SHAP 就退化成对单个标量的归因，省掉其余类别的计算
            return model(tensor)[:, target : target + 1].cpu().numpy()

    masker = shap.maskers.Image(mask_value, image.shape)
    explainer = shap.Explainer(predict, masker)
    explanation = explainer(image[None], max_evals=max_evals, batch_size=batch_size, silent=True)

    values = np.asarray(explanation.values[0])
    if values.ndim == 4:            # (H, W, C, n_outputs)
        values = values[..., 0]
    attribution = values.sum(axis=-1) if values.ndim == 3 else values

    positive = np.clip(attribution, 0.0, None)
    peak = positive.max()
    return positive / peak if peak > 0 else positive


def smooth(attribution: np.ndarray, radius: int = 2) -> np.ndarray:
    """方框平滑。像素级归因偏碎，直接取连通域会切出很多小块；
    轻微平滑后再阈值化，框更完整。radius=0 表示不平滑。"""
    if radius <= 0:
        return attribution
    size = 2 * radius + 1
    padded = np.pad(attribution, radius, mode="edge")
    # 用积分图做 O(1) 的方框均值，避免逐像素卷积
    integral = padded.cumsum(axis=0).cumsum(axis=1)
    integral = np.pad(integral, ((1, 0), (1, 0)), mode="constant")
    height, width = attribution.shape
    total = (
        integral[size:size + height, size:size + width]
        - integral[0:height, size:size + width]
        - integral[size:size + height, 0:width]
        + integral[0:height, 0:width]
    )
    smoothed = total / (size * size)
    peak = smoothed.max()
    return smoothed / peak if peak > 0 else smoothed


def batched_attribution(
    model: nn.Module,
    images: np.ndarray,
    targets: list[int],
    device: torch.device,
    progress: Callable[[int], None] | None = None,
    **kwargs,
) -> list[np.ndarray]:
    """逐张求归因。Partition explainer 本身是按单图工作的，这里只是循环 + 进度回调。"""
    results = []
    for index, (image, target) in enumerate(zip(images, targets, strict=True)):
        results.append(shap_attribution(model, image, target, device, **kwargs))
        if progress:
            progress(index)
    return results