"""
=============================================================================
CS329T · 第 3 讲数据与模型：带"角落水印"捷径的合成图像 + 一个小 CNN
=============================================================================

对应讲义：lectures/lecture03_deep_attribution.md §4–§7（代码 02、03、04 都导入本文件）

-----------------------------------------------------------------------------
【这个世界是怎么生成的 The recipe】
-----------------------------------------------------------------------------
第 1 讲 §1.4 讲过一个真实故事：一个从 X 光片判断癌症的模型准确率很高，归因却显示它最看重的像素落在
放射科医生留下的**笔迹**上 —— 一个"标签泄露"进了图像本身。这里用一个"知道真相"的小世界复现它：

    · 图像：16 × 16 灰度，背景是高斯噪声 N(0, 0.3²)。
    · 真正的信号：一根长 4–6 像素、宽 1 像素的短线，亮度 U(0.42, 0.78)，随机放在右下方 12 × 12 的区域里
      （第 4–15 行、第 4–15 列）。**类别 y = 1：竖线；y = 0：横线。** 这是我们"希望"模型学会的规律。
    · 捷径（shortcut）：左上角 3 × 3 的一个亮块（每个像素 +1.0），像片子上的笔迹 / 医院的标签。
      短线永远不会碰到这块区域，所以"归因落在这 9 个像素上"只可能来自水印。
    · 水印和标签的关系由 watermark 参数决定：
          "correlated"   y = 1 的图有 95% 带水印，y = 0 的图有 5% 带水印（历史数据：水印 ≈ 答案）
          "independent"  每张图 50% 带水印，与 y 无关（"修好了"的数据：水印没有信息）
          "none"         都不带水印（部署时：新医院、新扫描仪，没有这个标记）
          "flipped"      反过来：y = 0 的图 95% 带水印（最坏情况：标记的含义变了）
          "all"          每张图都带水印（用来做归因演示：水印在场时模型看不看它？）

两个模型，结构和初始化完全相同，只有训练数据不同：
    Model S（shortcut）：在 "correlated" 数据上训练 —— 它可以走捷径。
    Model C（control） ：在 "independent" 数据上训练 —— 它也见过水印，但水印没有用，只能学短线的方向。

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

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

torch.set_num_threads(1)

IMG = 16                    # 图像边长
NOISE_SD = 0.3              # 背景噪声的标准差
BAR_AMP = (0.42, 0.78)      # 短线亮度的范围
BAR_LEN = (4, 6)            # 短线长度的范围（含两端）
BAR_REGION = 4              # 短线只出现在第 4–15 行 / 列
WM_VALUE = 1.0              # 水印每个像素加多少亮度
P_WM = 0.95                 # "correlated" 时水印与标签一致的概率

# 水印区域的 bool 掩码：左上角 3 × 3，共 9 个像素，占全图 9 / 256 = 3.5%
PATCH = torch.zeros(IMG, IMG, dtype=torch.bool)
PATCH[0:3, 0:3] = True


def make_bars(n, rng, watermark="correlated", p_wm=P_WM):
    """
    生成 n 张"横线 / 竖线"图像，按 watermark 参数决定水印和标签的关系。

    【思路】先生成噪声背景，再逐张画一根短线（y = 0 横、y = 1 竖），最后在选中的图像左上角加水印。
        所有随机性都来自传入的 rng，同一个种子永远得到同一批图像。
    【参数】
        n          图像张数
        rng        np.random.Generator（np.random.default_rng(seed)）
        watermark  "correlated" / "independent" / "none" / "flipped" / "all"，含义见文件开头
        p_wm       correlated / flipped 时水印与标签一致（或相反）的概率
    【返回】dict：
        X    torch.float32，n × 1 × 16 × 16（"1"是通道数：灰度图只有一个通道）
        y    torch.float32，n 个 0/1 标签（BCE loss 要求 float）
        wm   numpy bool，n 个"是否带水印"
        bar  torch.bool，n × 16 × 16，短线所在的像素（"人类认为重要的区域"的标准答案）
    【边界】n = 0 也能运行（返回空张量）。短线完全落在第 4–15 行 / 列，永远不与水印重叠。
    【numpy 小知识】rng.integers(a, b) 返回 [a, b) 中的整数（不含 b）；rng.uniform(a, b) 返回 [a, b) 中的实数。
    """
    X = rng.normal(0.0, NOISE_SD, size=(n, 1, IMG, IMG)).astype(np.float32)
    y = rng.integers(0, 2, size=n)
    bar = np.zeros((n, IMG, IMG), dtype=bool)
    lo = BAR_REGION
    for i in range(n):
        L = int(rng.integers(BAR_LEN[0], BAR_LEN[1] + 1))
        if y[i] == 0:                                     # 横线：固定一行，占连续 L 列
            r = int(rng.integers(lo, IMG))
            c = int(rng.integers(lo, IMG - L + 1))
            bar[i, r, c:c + L] = True
        else:                                             # 竖线：固定一列，占连续 L 行
            c = int(rng.integers(lo, IMG))
            r = int(rng.integers(lo, IMG - L + 1))
            bar[i, r:r + L, c] = True
        X[i, 0][bar[i]] += rng.uniform(*BAR_AMP)          # 【Python 小知识】*BAR_AMP 把 (a, b) 拆成两个参数
    u = rng.random(n)
    if watermark == "correlated":
        wm = np.where(y == 1, u < p_wm, u < 1 - p_wm)
    elif watermark == "flipped":
        wm = np.where(y == 0, u < p_wm, u < 1 - p_wm)
    elif watermark == "independent":
        wm = u < 0.5
    elif watermark == "none":
        wm = np.zeros(n, dtype=bool)
    elif watermark == "all":
        wm = np.ones(n, dtype=bool)
    else:
        raise ValueError(f"未知的 watermark 模式：{watermark}")
    X[wm, 0, 0:3, 0:3] += WM_VALUE                        # 【numpy 小知识】bool 数组当下标：只改 wm 为 True 的图
    # 【PyTorch 小知识】torch.from_numpy 与 numpy 数组共享内存、不复制；torch.tensor 会复制一份。
    return {"X": torch.from_numpy(X), "y": torch.tensor(y, dtype=torch.float32), "wm": wm,
            "bar": torch.from_numpy(bar)}


def remove_watermark(X, wm):
    """
    返回 X 的一个副本：把带水印的图像的水印减掉（得到"同一张图、但没有水印"的反事实版本）。

    【思路】因为图像是我们自己合成的，知道水印加了多少，可以精确地"撤销"它 —— 现实数据里做不到，
        这正是合成世界的价值：可以直接测量"去掉水印，模型的输出会变多少"，用来检验归因的结论（代码 02 演示 4）。
    【参数】X：n × 1 × 16 × 16；wm：长度 n 的 numpy bool
    【返回】新的张量（不修改输入）
    """
    Z = X.clone()
    idx = torch.from_numpy(np.where(wm)[0])
    Z[idx, 0, 0:3, 0:3] -= WM_VALUE
    return Z


class SmallCNN(nn.Module):
    """
    一个很小的卷积网络（约 1.5k 个参数）：

        x (1×16×16) → conv1 3×3, 8 通道 → ReLU → maxpool 2 → (8×8×8)
                    → conv2 3×3, 16 通道 → ReLU → [h：16×8×8，"内部表示"] → maxpool 2 → (16×4×4)
                    → flatten (256) → fc → 1 个 logit（> 0 判为竖线 y = 1）

    【思路】把网络分成 body（输入 → conv2 的激活 h）和 head（h → logit）两段，代码 04 要在 h 这一层上
        计算每个单元（通道 × 位置，共 16 × 8 × 8 = 1024 个）的影响（conductance），并把某些单元"关掉"（ablation）：
        unit_mask 为 0 的单元，其激活被替换成 unit_fill 里的值（默认 0）。默认 mask 全 1，模型就是普通的 CNN。
        为什么用 flatten 而不是 global average pooling？flatten 保留了位置信息，模型才"有能力"学会
        "左上角有亮块"这条捷径 —— 真实网络通常也有足够的容量学会这种东西。
    【PyTorch 小知识】
        · nn.Module 是所有网络的基类；在 __init__ 里定义层，在 forward 里写前向计算，调用 model(x) 会执行 forward。
        · register_buffer：注册一个"不是参数、但属于模型状态"的张量（不参与训练，会随 state_dict 保存）。
        · F.max_pool2d(h, 2)：每个 2×2 的格子取最大值，图像边长减半。
    """

    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 8, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(8, 16, kernel_size=3, padding=1)
        self.fc = nn.Linear(16 * 4 * 4, 1)
        self.register_buffer("unit_mask", torch.ones(16, 8, 8))
        self.register_buffer("unit_fill", torch.zeros(16, 8, 8))

    def body(self, x):
        """输入 → conv2 之后（ReLU 之后、pooling 之前）的激活，形状 n × 16 × 8 × 8。"""
        h1 = F.max_pool2d(F.relu(self.conv1(x)), 2)
        return F.relu(self.conv2(h1))

    def head(self, h):
        """conv2 的激活 → logit（长度 n）。unit_mask 为 0 的单元被替换成 unit_fill（消融，代码 04）。"""
        h = h * self.unit_mask + (1.0 - self.unit_mask) * self.unit_fill
        return self.fc(F.max_pool2d(h, 2).flatten(1)).squeeze(1)

    def forward(self, x):
        return self.head(self.body(x))


def train_cnn(data, seed, epochs=10, lr=3e-3, batch=64):
    """
    用 Adam + binary cross-entropy 训练一个 SmallCNN。

    【思路】标准的小批量训练：每个 epoch 把训练集打乱，按 64 张一批更新参数。
        初始化用 torch.manual_seed(seed)，打乱顺序用 np.random.default_rng(seed)，所以同一个 seed 永远得到同一个模型。
        Model S 和 Model C 用**同一个** seed：两者的初始权重、更新次数完全一样，差别只来自数据。
    【参数】data：make_bars 返回的 dict；seed：整数；epochs、lr、batch：训练超参数
    【返回】训练好的 SmallCNN（已切换到 eval 模式）
    【边界】6000 张图、10 个 epoch 在单线程 CPU 上约 2 秒。没有早停、没有验证集调参 —— 演示用，够了。
    【PyTorch 小知识】
        · opt.zero_grad() 清掉上一步的梯度（PyTorch 默认会累加梯度）；loss.backward() 反向传播；opt.step() 更新参数。
        · F.binary_cross_entropy_with_logits 直接吃 logit，内部用数值稳定的方式算 sigmoid + log loss。
        · model.eval() 切换到推理模式（本模型没有 Dropout / BatchNorm，所以只是一个好习惯）。
    """
    torch.manual_seed(seed)
    model = SmallCNN()
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    rng = np.random.default_rng(seed)
    X, y = data["X"], data["y"]
    n = len(y)
    for _ in range(epochs):
        perm = torch.from_numpy(rng.permutation(n))
        for s in range(0, n, batch):
            idx = perm[s:s + batch]
            loss = F.binary_cross_entropy_with_logits(model(X[idx]), y[idx])
            opt.zero_grad()
            loss.backward()
            opt.step()
    return model.eval()


def accuracy(model, data):
    """准确率：logit > 0 判为 1。返回 Python float。"""
    with torch.no_grad():
        return ((model(data["X"]) > 0).float() == data["y"]).float().mean().item()


def build_world(n_train=6000, n_test=2000):
    """
    生成本讲所有图像实验共用的数据集和两个模型（代码 02、03、04 都调用它，所以三个脚本里的模型完全相同）。

    【返回】(sets, models)
        sets：dict，键为
            "train_S"   correlated（Model S 的训练集）
            "train_C"   independent（Model C 的训练集）
            "val"       correlated（"同一条数据管线"的验证集：排行榜用的就是它）
            "clean"     none（部署：没有水印）
            "flipped"   flipped（部署：水印含义反转）
            "all_wm"    all（每张都带水印，用来做归因）
        models：dict，"S" 和 "C"
    【边界】每个数据集用自己的种子（301–306），改其中一个不影响其他。
    """
    sets = {
        "train_S": make_bars(n_train, np.random.default_rng(301), "correlated"),
        "train_C": make_bars(n_train, np.random.default_rng(302), "independent"),
        "val": make_bars(n_test, np.random.default_rng(303), "correlated"),
        "clean": make_bars(n_test, np.random.default_rng(304), "none"),
        "flipped": make_bars(n_test, np.random.default_rng(305), "flipped"),
        "all_wm": make_bars(n_test, np.random.default_rng(306), "all"),
    }
    models = {"S": train_cnn(sets["train_S"], seed=0), "C": train_cnn(sets["train_C"], seed=0)}
    return sets, models


def evidence_fn(model, y):
    """
    返回"支持真实类别的证据"函数 f(z) = s · logit(z)，其中 s = +1（y = 1）或 −1（y = 0）。

    【思路】单个 logit 是"竖线 vs 横线"的对数几率。解释 y = 1 的图时问"为什么是竖线"，用 +logit；
        解释 y = 0 的图时问"为什么是横线"，用 −logit。这样所有图像上"归因为正 = 支持正确答案"。
    【参数】model：SmallCNN；y：长度 n 的 0/1 张量（与将来传入 f 的那批图一一对应）
    【返回】一个函数 f，输入 n 张图，返回长度 n 的张量
    【Python 小知识】在函数里定义并返回另一个函数（闭包 closure）：内层函数"记住"了外层的 model 和 s。
    """
    s = 2.0 * y - 1.0

    def f(z):
        return model(z) * s

    return f


if __name__ == "__main__":
    import time

    t0 = time.time()
    sets, models = build_world()
    print(f"数据 + 两个模型：{time.time() - t0:.1f} 秒")
    d = sets["val"]
    print("val 集：带水印比例 =", round(float(d["wm"].mean()), 3),
          "；只看水印猜标签的准确率 =", round(float((d["wm"] == d["y"].numpy().astype(bool)).mean()), 3))
    for k in ["val", "clean", "flipped"]:
        print(f"{k:<8} Model S = {accuracy(models['S'], sets[k]):.3f}   Model C = {accuracy(models['C'], sets[k]):.3f}")