"""
=============================================================================
CS230 · 第 21 讲工具模块：一个"知道答案在哪里"的图像世界、一个带 GAP 的小 CNN、五种归因方法
=============================================================================

对应讲义：lectures/lecture21_interpretability_and_wrapup.md §1.4、§2–§6

-----------------------------------------------------------------------------
【为什么需要它】
-----------------------------------------------------------------------------
可解释性方法最难的地方，是**没有标准答案**：一张显著图（saliency map）"看起来"对不对，全凭人眼。
所以本讲自己"画"一个世界，每张图的答案在哪里由我们的代码决定：

    · 32×32 的彩色图，**类别 = 一个"目标物体"的形状**（三角 / 十字 / 圆环 / 方块，4 类）；
      目标的位置、大小、旋转（±25°）、颜色都随机，颜色和类别无关。
    · 周围有 2 个小的**干扰物**（圆点、横条、竖条），它们和类别无关。
    · 背景是随机底色 + 渐变 + 像素噪声。
    · 可选：左下角盖一个 4×4 的**彩色"水印"**（代码 02 用它制造捷径：训练集里水印颜色和类别强相关）。
同时返回目标物体的**逐像素掩码和外接框**、干扰物的掩码、目标的中心和颜色——
有了它们，我们就能**量化**一个归因方法有没有"指对地方"，也能给线性探针（§6）提供要读的量。

模型是一个小 CNN，最后一层卷积之后接**全局平均池化（GAP）+ 一个线性层**。
这是 CAM（class activation map）要求的结构（讲义 §2.4）。

五种归因方法（都输出 (n, 32, 32) 的非负"重要性图"）：
    gradient_map          vanilla gradient / saliency map（Simonyan et al. 2014）：max_通道 |∂z_c/∂x|
    grad_x_input_map      gradient × input：把梯度乘上像素值本身，再在通道上求和
    integrated_gradients  Integrated Gradients（Sundararajan et al. 2017）：沿"基线 → 输入"的直线积分梯度
    occlusion_map         occlusion（Zeiler & Fergus 2014）：用灰块遮住一块，看类别分数掉多少
    cam_map / grad_cam_map  CAM（Zhou et al. 2016）与 Grad-CAM（Selvaraju et al. 2017）
外加一个**不看模型**的对照：edge_map（Sobel 边缘强度）。它用来回答："一个方法的图好看，会不会只是因为它像边缘检测？"

-----------------------------------------------------------------------------
【形状约定】PyTorch 的习惯：图像 (m, 3, H, W)，float32，像素大致在 [0, 1]；标签 y (m,) int64
-----------------------------------------------------------------------------
    归因图（attribution map）：(n, H, W) float64，非负，越大表示"越重要"
    框 box：(m, 4) int，[x0, y0, x1, y1]，**包含**两端（x 是列号、y 是行号）

依赖前几讲的模块（只读导入）：
    第 1 讲 console.py        pad、print_header（对齐表格、分隔标题）
    第 2 讲 text_plot.py      get_pyplot、figure_path（可选的 matplotlib 画图设置）
    第 9 讲 torch_utils.py    make_loader、fit、predict、count_parameters、setup_torch（训练循环模板）
    第 13 讲 shapes_world.py  shape_sdf（形状的有符号距离函数：本讲直接复用它来画形状）

【注释标签】【思路】【参数】【返回】【边界】【DL 联系】【numpy 小知识】【PyTorch 小知识】【Python 小知识】

直接运行本文件：画几张图的字符画，打印数据的基本统计，并检查五种归因方法的形状和 IG 的完整性（约 2 秒）：
    python3 interp_tools.py
"""

import os
import sys

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

HERE = os.path.dirname(os.path.abspath(__file__))
for _d in ("lecture01", "lecture02", "lecture09", "lecture13"):
    sys.path.insert(0, os.path.join(HERE, "..", _d))

from console import pad, print_header  # noqa: E402,F401  （第 1 讲；脚本从这里导入）
from shapes_world import shape_sdf  # noqa: E402            （第 13 讲：形状的有符号距离函数）
from text_plot import figure_path, get_pyplot  # noqa: E402,F401  （第 2 讲：画图设置）
from torch_utils import count_parameters, fit, make_loader, predict, setup_torch  # noqa: E402,F401  （第 9 讲）

S = 32                                              # 图像边长（像素）
CLASSES = ["triangle", "plus", "ring", "square"]    # 目标物体的 4 种形状；类别编号 = 下标
CLASS_CN = ["三角", "十字", "圆环", "方块"]
DISTRACTORS = ["disk", "hbar", "vbar"]              # 干扰物：和类别无关的小形状

# 水印（代码 02）：左下角 4×4 的纯色方块，4 种颜色与 4 个类别一一对应（红、绿、蓝、黄）
MARKER_Y0, MARKER_X0, MARKER_SIZE = 27, 1, 4
MARKER_COLORS = np.array([[0.95, 0.15, 0.15], [0.15, 0.85, 0.15], [0.15, 0.35, 0.95], [0.95, 0.85, 0.15]])
MARKER_CN = ["红", "绿", "蓝", "黄"]
# 物体不许进入的左下角区域（x ≤ 8 且 y ≥ 23）：所有图（有没有水印）都遵守，这样"有水印"和"没水印"的图只差水印本身
KEEP_OUT = (0, 23, 8, S - 1)                        # [x0, y0, x1, y1]


# =============================================================================
# 1. 世界：画图，并记下"答案在哪里"
# =============================================================================
def _far_color(bg, rng, min_dist=0.45):
    """抽一个和背景色 bg（长度 3）欧氏距离至少 min_dist 的颜色；抽 50 次都不行就用补色。返回 (3,)。"""
    for _ in range(50):
        c = rng.random(3)
        if np.linalg.norm(c - bg) >= min_dist:
            return c
    return 1.0 - bg


def _paint(img, name, cx, cy, R, theta, color):
    """
    在 img 上画一个形状（原地修改 img），返回这个形状的覆盖度 alpha。

    【思路】和第 13 讲 shapes_world.render_shapes 同样的画法：先把每个像素的坐标变到"以形状中心为原点、
            旋转 theta、以半径 R 为单位"的局部坐标 (u, v)，代入 shape_sdf 得到有符号距离 f（> 0 在里面），
            覆盖度 alpha = clip(0.5 + f·R, 0, 1)（边界上的像素在 0–1 之间，抗锯齿），
            新像素 = 旧像素 · (1 − alpha) + 颜色 · alpha。
    【参数】img：(3, S, S) float64；name：形状名；cx, cy：中心（像素）；R：半径（像素）；theta：旋转角（弧度）；
            color：(3,) 颜色
    【返回】alpha：(S, S)，0–1
    【numpy 小知识】pix[None, :] 是 (1, S) 的行向量、pix[:, None] 是 (S, 1) 的列向量，
            两者和标量相减后再相加时广播成 (S, S)：每个像素一个值。
    """
    pix = np.arange(S) + 0.5                              # 像素中心的坐标 0.5, 1.5, ..., 31.5
    dx, dy = pix[None, :] - cx, pix[:, None] - cy
    c, s = np.cos(theta), np.sin(theta)
    u, v = (c * dx + s * dy) / R, (-s * dx + c * dy) / R
    alpha = np.clip(0.5 + shape_sdf(name, u, v) * R, 0.0, 1.0)
    img[:] = img * (1 - alpha[None]) + color[:, None, None] * alpha[None]   # img[:] = ...：原地写回
    return alpha


def _boxes_overlap(a, b, margin=1.0):
    """两个框 [x0, y0, x1, y1] 在留出 margin 像素的间隔后是否重叠。"""
    return not (a[2] + margin < b[0] or b[2] + margin < a[0] or a[3] + margin < b[1] or b[3] + margin < a[1])


def _mask_box(mask):
    """一个布尔掩码的外接框 [x0, y0, x1, y1]（包含两端）。np.nonzero 返回 True 元素的行号数组、列号数组。"""
    rows, cols = np.nonzero(mask)
    return np.array([cols.min(), rows.min(), cols.max(), rows.max()])


def box_to_mask(box, tol=1):
    """
    把框变成 (S, S) 的布尔掩码，四边各放宽 tol 个像素（"指到框边上也算对"的容差）。

    【参数】box：[x0, y0, x1, y1]；tol：放宽的像素数
    【返回】(S, S) bool
    """
    m = np.zeros((S, S), dtype=bool)
    x0, y0 = max(box[0] - tol, 0), max(box[1] - tol, 0)
    x1, y1 = min(box[2] + tol, S - 1), min(box[3] + tol, S - 1)
    m[y0:y1 + 1, x0:x1 + 1] = True                     # 先行（y）后列（x）
    return m


def make_world(m, rng, marker="none", marker_agree=0.9, n_distract=2, noise=0.05):
    """
    生成 m 张图，以及每张图的"标准答案"。

    【思路】每张图依次：
            1. 背景：随机底色 + 随机方向的线性渐变；
            2. 目标物体：类别 y 均匀随机 → 形状 CLASSES[y]；半径 5.5–7.5 像素，中心随机（整个物体在图内、
               不进左下角的 KEEP_OUT 区域），旋转 ±25°，颜色随机（和背景差得足够远，和类别无关）；
            3. 干扰物：n_distract 个小形状（半径 2.6–3.6 像素），不和目标的框重叠、彼此不重叠、不进 KEEP_OUT；
            4. 像素噪声 N(0, noise²)；
            5. （可选）水印：在左下角盖一个 4×4 的纯色方块。像"照片来源"的角标一样，它是**事后盖上去的**，不带噪声。
    【参数】
        m            : 张数
        rng          : np.random.Generator（所有随机性都从它来，同一个种子画出同样的图）
        marker       : "none" 没有水印；"correlated" 水印颜色以概率 marker_agree 等于"类别对应的颜色"，
                       否则是另外 3 种颜色之一；"random" 水印颜色和类别无关（4 种颜色均匀随机）
        marker_agree : marker="correlated" 时水印和类别一致的概率
        n_distract   : 干扰物个数（放不下时会少一个，极少发生）
        noise        : 像素噪声的标准差
    【返回】dict：
        "X"      (m, 3, S, S) float32 图像
        "y"      (m,) int64 类别
        "mask"   (m, S, S) bool 目标物体的像素（alpha ≥ 0.5）
        "box"    (m, 4) int 目标物体的外接框 [x0, y0, x1, y1]
        "dmask"  (m, S, S) bool 干扰物的像素（所有干扰物合在一起）
        "center" (m, 2) float 目标中心 (cx, cy)，像素
        "color"  (m, 3) float 目标的颜色
        "marker" (m,) int 水印颜色编号（0–3），没有水印时是 −1
    【边界】物体的"大致框"用半径估计（形状最远伸到约 1.0 R），真实的框从掩码里量出来。
    【Python 小知识】for _ in range(50): ... else: ... —— for 循环的 else 分支只在循环"没有被 break"时执行。
            这里用不到它；我们用 placed 标志来判断有没有放下。
    """
    X = np.zeros((m, 3, S, S), dtype=np.float32)
    y = rng.integers(0, len(CLASSES), m).astype(np.int64)
    mask = np.zeros((m, S, S), dtype=bool)
    dmask = np.zeros((m, S, S), dtype=bool)
    box = np.zeros((m, 4), dtype=np.int64)
    center = np.zeros((m, 2))
    color = np.zeros((m, 3))
    marker_id = -np.ones(m, dtype=np.int64)
    # 水印的颜色用一个**单独的**随机数生成器抽：这样 marker="correlated" 和 marker="random" 用同一个种子时，
    # 物体、背景、噪声逐像素相同，只有水印颜色不同（代码 02 演示 4 要用这一点）。marker="none" 时不抽，世界和以前一样。
    mrng = np.random.default_rng(int(rng.integers(2 ** 31))) if marker != "none" else None
    pix = np.arange(S) + 0.5
    for i in range(m):
        # 1. 背景
        bg = rng.random(3) * 0.8 + 0.1
        ang = rng.uniform(0, 2 * np.pi)
        ramp = (pix[None, :] / S - 0.5) * np.cos(ang) + (pix[:, None] / S - 0.5) * np.sin(ang)
        img = bg[:, None, None] + rng.uniform(0, 0.25) * ramp[None]
        # 2. 目标物体：重抽中心，直到大致框不进 KEEP_OUT
        R = rng.uniform(5.5, 7.5)
        while True:
            cx, cy = rng.uniform(R + 1, S - R - 1, 2)
            rough = [cx - R - 0.5, cy - R - 0.5, cx + R + 0.5, cy + R + 0.5]
            if not _boxes_overlap(rough, KEEP_OUT, margin=0.0):
                break
        col = _far_color(bg, rng)
        alpha = _paint(img, CLASSES[y[i]], cx, cy, R, np.deg2rad(rng.uniform(-25, 25)), col)
        mask[i] = alpha >= 0.5
        box[i] = _mask_box(mask[i])
        center[i], color[i] = (cx, cy), col
        # 3. 干扰物
        placed = [box[i]]
        for _ in range(n_distract):
            for _ in range(50):
                r = rng.uniform(2.6, 3.6)
                dx_, dy_ = rng.uniform(r + 0.5, S - r - 0.5, 2)
                rough = [dx_ - r - 0.5, dy_ - r - 0.5, dx_ + r + 0.5, dy_ + r + 0.5]
                if any(_boxes_overlap(rough, b) for b in placed) or _boxes_overlap(rough, KEEP_OUT, margin=0.0):
                    continue
                name = DISTRACTORS[int(rng.integers(0, len(DISTRACTORS)))]
                a2 = _paint(img, name, dx_, dy_, r, np.deg2rad(rng.uniform(-25, 25)), _far_color(bg, rng))
                dmask[i] |= a2 >= 0.5                       # |= ：逐元素"或"，把这个干扰物并进去
                placed.append(rough)
                break
        # 4. 噪声
        img = img + rng.normal(0.0, noise, img.shape)
        # 5. 水印
        if marker == "correlated":
            k = y[i] if mrng.random() < marker_agree else (y[i] + mrng.integers(1, 4)) % 4   # 其余 3 种颜色之一
            marker_id[i] = k
        elif marker == "random":
            marker_id[i] = mrng.integers(0, 4)
        elif marker != "none":
            raise ValueError(f"unknown marker mode: {marker}")
        X[i] = img
    if marker != "none":
        X = paint_marker(X, marker_id)
    return {"X": X, "y": y, "mask": mask, "box": box, "dmask": dmask, "center": center, "color": color,
            "marker": marker_id}


def paint_marker(X, marker_id):
    """
    在每张图的左下角盖上第 marker_id[i] 种颜色的水印（返回新数组，不改原图）。

    【思路】水印是"事后盖上去"的纯色方块，所以换颜色只要覆盖那 16 个像素——这让代码 02 的**干预实验**
            （把水印换成别的颜色，看预测跟不跟着变）非常干净：除了水印，图上一个像素都不变。
    【参数】X：(m, 3, S, S)；marker_id：(m,) int，0–3（−1 表示这张图不盖）
    【返回】(m, 3, S, S) float32
    """
    X = X.copy()
    ys, xs = slice(MARKER_Y0, MARKER_Y0 + MARKER_SIZE), slice(MARKER_X0, MARKER_X0 + MARKER_SIZE)
    for i, k in enumerate(marker_id):
        if k >= 0:
            X[i, :, ys, xs] = MARKER_COLORS[k][:, None, None]
    return X.astype(np.float32)


def marker_region_mask():
    """水印所在的 4×4 区域（(S, S) bool），用来统计"归因有多少落在水印上"。"""
    m = np.zeros((S, S), dtype=bool)
    m[MARKER_Y0:MARKER_Y0 + MARKER_SIZE, MARKER_X0:MARKER_X0 + MARKER_SIZE] = True
    return m


# =============================================================================
# 2. 模型：GAP 结尾的小 CNN
# =============================================================================
class GapCNN(nn.Module):
    """
    一个以"全局平均池化 + 线性层"结尾的小 CNN。

    【结构】（输入 3 × 32 × 32，像素 0–1）
            标准化（固定的按通道均值 / 标准差，写在模型里：这样归因直接是对**原始像素**求的）
            conv1  3×3，12 通道 → BN → ReLU → maxpool 2   → 12 × 16 × 16
            conv2  3×3，24 通道 → BN → ReLU → maxpool 2   → 24 × 8 × 8
            conv3  3×3，32 通道 → BN → ReLU               → 32 × 8 × 8
            conv4  3×3，48 通道 → BN → ReLU               → A：48 × 8 × 8（"最后一层特征图"）
            GAP    每个通道在 8×8 上取平均                  → 48 个数
            fc     线性层 48 → 4                           → 4 个 logits
    【思路】CAM 要求"最后一层卷积 → GAP → 线性层"：这时 logit_c = Σ_k w_ck · mean_ij A_k(i, j) + b_c
            = mean_ij [Σ_k w_ck A_k(i, j)] + b_c —— 方括号里那张 8×8 的图就是 CAM，它把 logit 按位置拆开了（讲义 §2.4）。
            BatchNorm 让这个小网络训练得又快又稳（不加 BN 时我在草稿里试过，前几个 epoch 常常卡在平台上）。
    【参数】mean、std：(3,) 按通道的均值和标准差（用训练集算）；widths：四个卷积层的通道数；n_classes：类别数
    【PyTorch 小知识】self.register_buffer(名字, 张量)：登记一个"不是参数"的张量——它跟着模型保存、搬设备，
            但不参与梯度下降（model.parameters() 里没有它）。BN 的 running_mean 也是 buffer。
    """

    def __init__(self, mean, std, widths=(12, 24, 32, 48), n_classes=4):
        super().__init__()
        self.register_buffer("mean", torch.as_tensor(np.asarray(mean), dtype=torch.float32).view(1, 3, 1, 1))
        self.register_buffer("std", torch.as_tensor(np.asarray(std), dtype=torch.float32).view(1, 3, 1, 1))
        c1, c2, c3, c4 = widths

        def block(cin, cout, pool):
            layers = [nn.Conv2d(cin, cout, 3, padding=1), nn.BatchNorm2d(cout), nn.ReLU()]
            if pool:
                layers.append(nn.MaxPool2d(2))
            return nn.Sequential(*layers)

        # nn.ModuleList：一个"装着若干层"的列表，里面的层的参数会被自动登记（普通 Python 列表不会！）
        self.blocks = nn.ModuleList([block(3, c1, True), block(c1, c2, True), block(c2, c3, False), block(c3, c4, False)])
        self.fc = nn.Linear(c4, n_classes)

    def features(self, x, upto=4):
        """
        前向到第 upto 个卷积块为止，返回那一块的输出（ReLU / 池化之后）。

        【参数】x：(n, 3, S, S) 原始像素；upto：1–4
        【返回】(n, C, h, w) 的特征图；upto = 4 时就是 CAM 用的 A（48 × 8 × 8）
        """
        h = (x - self.mean) / self.std
        for blk in self.blocks[:upto]:
            h = blk(h)
        return h

    def forward(self, x):
        """x：(n, 3, S, S) → logits：(n, n_classes)。A.mean(dim=(2, 3)) 就是 GAP。"""
        return self.fc(self.features(x).mean(dim=(2, 3)))


def channel_stats(X):
    """训练集按通道的均值、标准差：(3,)、(3,)。axis=(0, 2, 3) 表示对"样本、高、宽"三个维度一起求。"""
    return X.mean(axis=(0, 2, 3)), X.std(axis=(0, 2, 3))


def train_world_model(data, seed=0, epochs=10, lr=3e-3, batch_size=64, widths=(12, 24, 32, 48)):
    """
    在一份 make_world 数据上训练一个 GapCNN（Adam + 余弦学习率衰减，第 9 讲 torch_utils.fit）。

    【思路】种子在**建模型之前**设（层的初始化用全局随机数，第 9 讲 §4.6）；DataLoader 的打乱用同一个 seed 的
            torch.Generator。所以同一个 (data, seed) 每次训练出逐位相同的模型——代码 01、02、03 都调用它，
            用同样的数据和种子时得到的是**同一个**模型。
    【参数】data：make_world 的返回值；seed：种子；epochs、lr、batch_size：训练超参数；widths：通道数
    【返回】(model, history)：训练好的模型（已经切到 eval 模式）；fit 的历史记录
    【DL 联系】余弦衰减（CosineAnnealingLR）是第 8 讲学习率调度的一种：学习率从 lr 平滑地降到 0。
    """
    mean, std = channel_stats(data["X"])
    torch.manual_seed(seed)
    model = GapCNN(mean, std, widths)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)
    loader = make_loader(data["X"], data["y"], batch_size, shuffle=True, seed=seed)
    hist, _ = fit(model, opt, loader, None, epochs=epochs, scheduler=sched)
    model.eval()
    return model, hist


def accuracy(model, X, y):
    """分类准确率（Python float）。X：(m, 3, S, S)；y：(m,)。"""
    return float(np.mean(predict(model, X).argmax(dim=1).numpy() == y))


# =============================================================================
# 3. 归因方法：每个都回答"哪些像素让类别 c 的分数这么高"
# =============================================================================
def _as_tensor(X):
    """numpy 数组（或张量）→ float32 张量。np.array(...) 会复制一份：广播出来的只读数组也能安全地交给 PyTorch。"""
    if isinstance(X, torch.Tensor):
        return X.to(torch.float32)
    return torch.from_numpy(np.array(X, dtype=np.float32))


def class_logits(model, X, cls):
    """
    第 i 张图的第 cls[i] 个 logit（不经过 softmax 的类别分数）。

    【参数】X：(n, 3, S, S) numpy 或张量；cls：(n,) int
    【返回】(n,) numpy float64
    【DL 联系】归因一般解释 logit 而不是 softmax 概率：概率接近 1 时会"饱和"，梯度几乎为 0（第 2 讲的 sigmoid 饱和）。
    """
    with torch.no_grad():
        z = model(_as_tensor(X))
    return z[torch.arange(len(cls)), torch.as_tensor(cls)].double().numpy()


def input_gradient(model, X, cls):
    """
    ∂ z_c / ∂ x：类别 c 的 logit 对每个输入像素的梯度。

    【思路】和第 5 讲 §4.5（风格迁移、对抗样本）一样："固定参数，对**输入**求梯度"。
            一批图一起算：把 n 张图各自的 z_{cls[i]} 加起来再 backward——因为第 i 张图的 logit 只依赖第 i 张图，
            和的梯度在第 i 张图上就等于它自己那个 logit 的梯度。
    【参数】model：eval 模式的模型；X：(n, 3, S, S)；cls：(n,) int
    【返回】(n, 3, S, S) numpy float64
    【PyTorch 小知识】x.requires_grad_(True)：告诉 autograd"要对这个张量求导"。默认只有参数要求梯度，输入不要求。
            .detach()：从计算图上取下来，之后再 .numpy()。
    """
    x = _as_tensor(X).clone().requires_grad_(True)
    z = model(x)
    z[torch.arange(len(cls)), torch.as_tensor(cls)].sum().backward()
    return x.grad.detach().double().numpy()


def gradient_map(model, X, cls):
    """
    vanilla gradient / saliency map（Simonyan, Vedaldi & Zisserman 2014）：每个像素取 3 个通道里 |梯度| 的最大值。

    【思路】一阶泰勒展开：z_c(x + δ) ≈ z_c(x) + ∇z_c · δ。|∂z_c/∂x_ij| 大 = "这个像素稍微一变，分数就变得多"。
            梯度的正负号对单个像素意义不大（变亮还是变暗都可能让分数变），所以取绝对值；通道取最大是原论文的做法。
    【返回】(n, S, S) 非负
    【边界】它回答的是"局部敏感度"，不是"贡献"：一个早已饱和、对分数贡献很大的特征，梯度可能很小（讲义 §2.5）。
    """
    return np.abs(input_gradient(model, X, cls)).max(axis=1)


def grad_x_input_map(model, X, cls):
    """
    gradient × input：Σ_通道 x · ∂z_c/∂x，只保留正的部分（"支持类别 c 的证据"）。

    【思路】梯度说"多敏感"，乘上像素值说"这个像素实际放了多少进来"。对线性模型 z = w·x + b，
            w_i · x_i 正好是第 i 个输入对 z 的贡献（相对于全 0 输入）。
    【返回】(n, S, S) 非负
    【边界】乘上了输入本身，所以图里天然带着"输入长什么样"的结构——这正是讲义 §3.2 的 sanity check 要警惕的。
    """
    g = input_gradient(model, X, cls)
    return np.maximum((np.asarray(X, dtype=np.float64) * g).sum(axis=1), 0.0)


def integrated_gradients(model, X, cls, baseline, steps=32, batch=256):
    """
    Integrated Gradients（Sundararajan, Taly & Yan 2017）。

    【思路】选一个"什么都没有"的基线 x'（比如全黑的图），沿直线 x' + α(x − x')（α 从 0 到 1）一路走过去，
            把每一点的梯度平均起来，再乘上 (x − x')：
                IG_i(x) = (x_i − x'_i) · ∫₀¹ ∂z_c(x' + α(x − x'))/∂x_i dα
            积分用中点法则近似：α_k = (k − 0.5)/steps，k = 1..steps。
            **完整性（completeness）**：Σ_i IG_i = z_c(x) − z_c(x')（微积分基本定理，对路径积分成立）。
            所以 IG 把"分数比基线高了多少"不多不少地分给了每个像素；steps 越多，近似越准（代码 01 演示 3 检验）。
    【参数】
        model    : eval 模式的模型
        X        : (n, 3, S, S) 输入
        cls      : (n,) 要解释的类别
        baseline : (n, 3, S, S) 或 (3, S, S) 的基线图像（全黑、平均色、模糊图……）
        steps    : 积分的步数
        batch    : 一次送进网络的图片数（只影响速度和内存）
    【返回】(ig, z_x, z_base)：ig (n, 3, S, S) float64 带符号的归因；z_x、z_base (n,) 输入和基线的 logit
    【边界】结果依赖基线的选择：基线里已经有的东西（比如全黑基线下的黑色像素）得不到归因（讲义 §2.5）。
    """
    X = np.asarray(X, dtype=np.float64)
    base = np.broadcast_to(np.asarray(baseline, dtype=np.float64), X.shape)   # broadcast_to：把 (3,S,S) 当作 (n,3,S,S) 用，不复制
    diff = X - base
    total = np.zeros_like(X)
    alphas = (np.arange(1, steps + 1) - 0.5) / steps
    n = len(X)
    for a in alphas:
        pts = base + a * diff                                 # 路径上的一点：(n, 3, S, S)
        for s0 in range(0, n, batch):                         # 分批：一次最多 batch 张
            sl = slice(s0, s0 + batch)
            total[sl] += input_gradient(model, pts[sl], cls[sl])
    ig = diff * total / steps
    return ig, class_logits(model, X, cls), class_logits(model, base, cls)


def ig_map(model, X, cls, baseline, steps=32):
    """IG 的重要性图：Σ_通道 IG，只保留正的部分。返回 (n, S, S)。"""
    ig, _, _ = integrated_gradients(model, X, cls, baseline, steps)
    return np.maximum(ig.sum(axis=1), 0.0)


def occlusion_map(model, X, cls, fill, patch=8, stride=4, batch=512):
    """
    occlusion（遮挡敏感度，Zeiler & Fergus 2014）：用一块灰色方块遮住图的一部分，看类别分数掉多少。

    【思路】对每个遮挡位置 p（patch × patch 的方块，步长 stride），算 drop_p = z_c(x) − z_c(x 遮住 p)。
            一个像素的重要性 = 所有盖住它的方块的 drop 的平均，再只保留正的部分（"遮住它，分数会掉"）。
            它**只用前向传播**，不需要梯度，也不需要知道网络内部的结构——任何黑盒模型都能用。
    【参数】X：(n, 3, S, S)；cls：(n,)；fill：(3,) 遮挡块的颜色（用训练集的平均色）；patch、stride：方块大小和步长；
            batch：一次前向的图片数
    【返回】(n, S, S) 非负
    【边界】代价：每张图要做 ((S − patch)/stride + 1)² 次前向（这里 7 × 7 = 49 次）；分辨率受 patch 和 stride 限制；
            方块大小要和物体差不多——太小遮不住、太大一次遮掉好几样东西。灰块本身也是"不自然"的输入（分布外）。
    """
    X = np.asarray(X, dtype=np.float32)
    n = len(X)
    base = class_logits(model, X, cls)
    starts = list(range(0, S - patch + 1, stride))
    heat = np.zeros((n, S, S))
    count = np.zeros((S, S))
    fill = np.asarray(fill, dtype=np.float32)[:, None, None]
    for y0 in starts:
        for x0 in starts:
            Xo = X.copy()
            Xo[:, :, y0:y0 + patch, x0:x0 + patch] = fill
            drop = np.concatenate([base[s:s + batch] - class_logits(model, Xo[s:s + batch], cls[s:s + batch])
                                   for s in range(0, n, batch)])
            heat[:, y0:y0 + patch, x0:x0 + patch] += drop[:, None, None]   # 这块里的每个像素都记上这次的 drop
            count[y0:y0 + patch, x0:x0 + patch] += 1
    return np.maximum(heat / np.maximum(count, 1), 0.0)


def _upsample(M):
    """(n, h, w) → (n, S, S)，双线性插值。F.interpolate 要 (n, 通道, h, w)，所以先加一个通道维再去掉。"""
    t = torch.as_tensor(M, dtype=torch.float64)[:, None]
    return F.interpolate(t, size=(S, S), mode="bilinear", align_corners=False)[:, 0].numpy()


def cam_map(model, X, cls):
    """
    CAM（class activation map，Zhou et al. 2016）：用最后一个线性层的权重，把最后一层特征图加权求和。

    【思路】logit_c = mean_ij [Σ_k w_ck A_k(i, j)] + b_c，所以 M_c(i, j) = Σ_k w_ck A_k(i, j) 告诉我们
            "logit 的哪一部分来自位置 (i, j)"。M_c 是 8 × 8 的，双线性插值放大到 32 × 32，再只保留正的部分。
            它**只需要一次前向**，而且严格对应模型的计算——代价是模型必须是"卷积 → GAP → 线性层"的结构。
    【参数】X：(n, 3, S, S)；cls：(n,)
    【返回】(n, S, S) 非负
    【numpy 小知识】np.einsum("nk,nkhw->nhw", W, A)：对每个 n，把 W[n, k] 和 A[n, k, h, w] 相乘后在 k 上求和。
            einsum 用下标字符串写出"哪些维度相乘、哪些求和"，比一串 reshape + 矩阵乘法好读。
    """
    with torch.no_grad():
        A = model.features(_as_tensor(X)).double().numpy()             # (n, 48, 8, 8)
    W = model.fc.weight.detach().double().numpy()[np.asarray(cls)]     # (n, 48)：每张图取它要解释的那一类的权重
    return np.maximum(_upsample(np.einsum("nk,nkhw->nhw", W, A)), 0.0)


def grad_cam_map(model, X, cls):
    """
    Grad-CAM（Selvaraju et al. 2017）：用"logit 对特征图的梯度的平均值"代替 CAM 里的权重 w_ck。

    【思路】α_k = mean_ij ∂z_c/∂A_k(i, j)，Grad-CAM = ReLU(Σ_k α_k A_k)。好处是不要求 GAP 结构，任何 CNN 都能用。
            在 GAP + 线性层的网络上，∂z_c/∂A_k(i, j) = w_ck / (h·w)，于是 Grad-CAM = CAM / (h·w)——
            两者只差一个常数（代码 01 演示 2 数值验证）。
    【返回】(n, S, S) 非负
    【PyTorch 小知识】A.retain_grad()：A 是中间结果（不是叶子张量），默认 backward 之后不保留它的 .grad，要显式要求。
    """
    x = _as_tensor(X)
    A = model.features(x)
    A.retain_grad()
    z = model.fc(A.mean(dim=(2, 3)))
    z[torch.arange(len(cls)), torch.as_tensor(cls)].sum().backward()
    alpha = A.grad.mean(dim=(2, 3), keepdim=True)                      # (n, 48, 1, 1)
    M = (alpha * A).sum(dim=1).detach().double().numpy()               # (n, 8, 8)
    return np.maximum(_upsample(M), 0.0)


def edge_map(X):
    """
    **不看模型**的对照：每个通道做 Sobel 边缘检测，梯度幅值在通道上求和。

    【思路】Adebayo et al.（2018）的提醒：有些显著图看起来很"对"，其实和一个边缘检测器差不多——
            而边缘检测器根本不知道模型学到了什么。把它和真正的归因方法放在同一张表里，才知道"好看"值多少钱。
    【参数】X：(n, 3, S, S)
    【返回】(n, S, S) 非负
    【DL 联系】Sobel 核就是第 12 讲 §3.4 的那个 3×3 核；这里用 F.conv2d 一次处理所有通道（groups=3：每个通道各卷各的）。
    """
    kx = torch.tensor([[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]], dtype=torch.float64)
    k = torch.stack([kx, kx.T])[:, None].repeat(3, 1, 1, 1)           # (6, 1, 3, 3)：每个通道一对 (x, y) 核
    x = torch.as_tensor(np.asarray(X), dtype=torch.float64)
    g = F.conv2d(F.pad(x, (1, 1, 1, 1), mode="replicate"), k, groups=3)   # (n, 6, S, S)
    g = g.view(len(x), 3, 2, S, S)
    return torch.sqrt((g ** 2).sum(dim=2)).sum(dim=1).numpy()


def all_maps(model, X, cls, fill, ig_steps=32):
    """
    一次算出所有方法的重要性图（IG 用全黑基线）。

    【返回】dict：方法名 → (n, S, S)。键的顺序就是表格里的行序。
    【Python 小知识】Python 3.7 起，dict 会记住键插入的顺序，所以遍历时的顺序是固定的。
    """
    return {
        "gradient": gradient_map(model, X, cls),
        "grad×input": grad_x_input_map(model, X, cls),
        "IG（黑基线）": ig_map(model, X, cls, np.zeros((3, S, S)), ig_steps),
        "occlusion": occlusion_map(model, X, cls, fill),
        "CAM": cam_map(model, X, cls),
        "edge（不看模型）": edge_map(X),
    }


# =============================================================================
# 4. 指标：归因"指对地方"了吗？两张图有多像？
# =============================================================================
def pointing_hits(maps, boxes, tol=1):
    """
    pointing game（Zhang et al. 2016）：重要性最大的那个像素，是否落在目标物体的框里（四边放宽 tol 像素）。

    【参数】maps：(n, S, S)；boxes：(n, 4)；tol：容差
    【返回】(n,) bool
    【边界】一张图全为 0（比如 CAM 的正部分处处为 0）时，argmax 返回第 0 个像素（左上角）——算作没指中，除非框恰好在那儿。
    【numpy 小知识】np.unravel_index(扁平下标, 形状)：把"拉平后第几个"换回 (行, 列)。
    """
    hits = np.zeros(len(maps), dtype=bool)
    for i, (M, b) in enumerate(zip(maps, boxes)):
        r, c = np.unravel_index(np.argmax(M), M.shape)
        hits[i] = (b[0] - tol <= c <= b[2] + tol) and (b[1] - tol <= r <= b[3] + tol)
    return hits


def mass_share(maps, region_masks):
    """
    每张图的重要性有多大比例落在给定区域里：Σ_区域 M / Σ M。

    【参数】maps：(n, S, S) 非负；region_masks：(n, S, S) 或 (S, S) bool
    【返回】(n,)；一张图的重要性全为 0 时记为 0
    """
    tot = maps.sum(axis=(1, 2))
    inside = (maps * region_masks).sum(axis=(1, 2))
    return np.where(tot > 0, inside / np.maximum(tot, 1e-300), 0.0)


def rank_average(a):
    """
    把一维数组换成"名次"（0 起），**相同的值取平均名次**。例：[10, 30, 20, 20] → [0, 3, 1.5, 1.5]。

    【思路】先排序；相同值在排好序的数组里连成一段，这一段的名次取平均。
            为什么要处理并列？CAM、occlusion 只保留正的部分，图里常有一大片 0。如果并列的 0 按下标排名次，
            两张图只要"0 在同一片区域"就会显得高度相关——那是假象。
    【参数】a：(N,)
    【返回】(N,) float64
    【numpy 小知识】np.argsort(a, kind="mergesort")：稳定排序，返回"排序后第 k 个元素在原数组里的下标"；
            np.flatnonzero(布尔数组)：True 的下标；np.repeat(值, 次数)：每个值重复对应的次数。
    """
    order = np.argsort(a, kind="mergesort")
    s = a[order]
    starts = np.flatnonzero(np.r_[True, s[1:] != s[:-1], True])   # 每一段的起点，最后加上终点 N
    counts = np.diff(starts)
    avg = starts[:-1] + (counts - 1) / 2.0                          # 这一段的平均名次
    ranks = np.empty(len(a))
    ranks[order] = np.repeat(avg, counts)
    return ranks


def spearman(a, b):
    """
    Spearman 秩相关：先把两张图各自换成名次，再算名次之间的 Pearson 相关。只看"排序像不像"，不看数值大小。

    【参数】a、b：同形状的数组（会被拉平）
    【返回】float，−1 到 1；有一张图是常数时返回 0
    """
    ra, rb = rank_average(np.ravel(a)), rank_average(np.ravel(b))
    ra, rb = ra - ra.mean(), rb - rb.mean()
    den = np.sqrt((ra ** 2).sum() * (rb ** 2).sum())
    return float((ra * rb).sum() / den) if den > 0 else 0.0


# =============================================================================
# 5. 字符画
# =============================================================================
def ascii_map(M, size=16, chars=" .:-=+*#%@", scale=None):
    """
    把一张 (S, S) 的图缩小到 size × size（平均池化）再画成字符，每格 2 个字符。

    【参数】M：(S, S) 非负；size：缩小后的边长（S 要能被它整除）；scale：归一化用的最大值（默认用 M 自己的）
    【返回】list[str]，size 行
    【numpy 小知识】M.reshape(16, 2, 16, 2).mean(axis=(1, 3))：把 32×32 切成 16×16 个 2×2 小块，每块取平均。
    """
    k = S // size
    m = np.asarray(M, dtype=np.float64).reshape(size, k, size, k).mean(axis=(1, 3))
    top = float(m.max()) if scale is None else scale
    idx = np.clip((m / max(top, 1e-12) * (len(chars) - 1)).round().astype(int), 0, len(chars) - 1)
    return ["".join(chars[v] * 2 for v in row) for row in idx]


def ascii_image(img, size=16):
    """一张彩色图 (3, S, S) 画成字符：先变成亮度（三个通道平均），再缩小。返回 list[str]。"""
    return ascii_map(np.clip(np.asarray(img).mean(axis=0), 0, 1), size, scale=1.0)


def side_by_side(blocks, titles, gap=3):
    """
    把几块字符画并排：返回可以逐行 print 的 list[str]。

    【Python 小知识】zip(*blocks)：先取每块的第 0 行凑成一组，再取第 1 行……（要求每块行数相同）。
    """
    widths = [max(len(line) for line in b) for b in blocks]
    sep = " " * gap
    lines = [sep.join(pad(t, w) for t, w in zip(titles, widths))]
    for rows in zip(*blocks):
        lines.append(sep.join(pad(r, w) for r, w in zip(rows, widths)))
    return lines


# =============================================================================
# 自检
# =============================================================================
if __name__ == "__main__":
    setup_torch(1)
    d = make_world(6, np.random.default_rng(0))
    print_header("make_world：6 张图（字符画是亮度；第二行是目标物体的掩码）")
    print("\n".join(side_by_side([ascii_image(x, 16) for x in d["X"][:4]],
                                [f"#{i} {CLASS_CN[d['y'][i]]}" for i in range(4)])))
    print("\n".join(side_by_side([ascii_map(m.astype(float), 16) for m in d["mask"][:4]],
                                ["目标掩码"] * 4)))
    big = make_world(2000, np.random.default_rng(1))
    area = big["mask"].mean()
    box_area = np.mean([box_to_mask(b, 0).mean() for b in big["box"]])
    print(f"\n2000 张图：每类 {np.bincount(big['y']).tolist()} 张；目标物体平均占 {area:.1%} 的像素，外接框占 {box_area:.1%}；"
          f"干扰物占 {big['dmask'].mean():.1%}")
    dm = make_world(8, np.random.default_rng(2), marker="correlated")
    print(f"marker='correlated'：类别 {dm['y'].tolist()}，水印颜色 {dm['marker'].tolist()}")

    # 用一个随机初始化的小模型检查五种方法的形状，以及 IG 的完整性
    torch.manual_seed(0)
    model = GapCNN(*channel_stats(big["X"])).eval()
    X, cls = d["X"], d["y"]
    for name, M in all_maps(model, X, cls, fill=channel_stats(big["X"])[0]).items():
        print(f"  {pad(name, 18)} 形状 {M.shape}，最小值 {M.min():.2e}（应 ≥ 0）")
    ig, zx, zb = integrated_gradients(model, X, cls, np.zeros((3, S, S)), steps=256)
    err = np.abs(ig.sum(axis=(1, 2, 3)) - (zx - zb)).max()
    print(f"IG 完整性（随机模型，256 步）：max |Σ IG − (z(x) − z(x'))| = {err:.2e}")
    g_cam, cam = grad_cam_map(model, X, cls), cam_map(model, X, cls)
    unit = lambda M: M / np.maximum(M.max(axis=(1, 2), keepdims=True), 1e-12)     # 每张图除以自己的最大值
    print(f"Grad-CAM 与 CAM 的最大差（各自除以最大值后）：{np.abs(unit(g_cam) - unit(cam)).max():.2e}")
    print(f"rank_average([10, 30, 20, 20]) = {rank_average(np.array([10.0, 30.0, 20.0, 20.0])).tolist()}")