"""
=============================================================================
CS230 · 第 13 讲工具模块：按步数训练、冻结 / 微调、数据增强、预训练一个"源模型"
=============================================================================

对应讲义：lectures/lecture13_deep_cnn_architectures.md（§3 plain vs residual、§7 迁移学习、§8 数据增强）

-----------------------------------------------------------------------------
【这个模块里有什么】
-----------------------------------------------------------------------------
    train_steps        按"步数"训练（而不是按 epoch）：小数据集和大数据集用同样多的梯度更新步数时才好比较；
                       支持数据增强、只训练一部分参数、让冻结的 BN 层保持 eval 模式
    accuracy、features  评估准确率；把图像送进（冻结的）网络前半段，取出特征
    TransferCNN        迁移学习实验用的小 CNN：block1 → block2 → block3 → head，方便"冻结前面几块"
    pretrain_source    在"大"的源任务（16 种形状 × 600 张）上预训练一个 TransferCNN —— 我们的"ImageNet 预训练模型"
    flip_lr、translate、color_jitter、pca_color、compose   数据增强（都在一个 batch 上做，用自己的随机数生成器）

依赖第 9 讲的 torch_utils.py（predict、count_parameters、setup_torch），只读导入。

-----------------------------------------------------------------------------
【训练循环还是那五行】
-----------------------------------------------------------------------------
train_steps 的循环体和第 9 讲 torch_utils.train_one_epoch 一样：
    logits = model(xb); loss = loss_fn(logits, yb); optimizer.zero_grad(); loss.backward(); optimizer.step()
多出来的只有三件事：(1) 每取一个 batch 先做数据增强；(2) 只把要训练的参数交给优化器；
(3) model.train() 之后，把"冻结"的子模块切回 eval()，让它们的 BatchNorm 继续用预训练时的 running 统计量。
第 (3) 条是迁移学习里最容易踩的坑（讲义 §7.5），代码 05 演示 3 专门演示忘了它会怎样。

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

直接运行本文件做一个自检（约 25 秒）：预训练两次源模型，确认结果逐位相同；打印每种增强对一个 batch 的作用。
    python3 transfer_kit.py
"""

import os
import sys

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

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture09"))
from torch_utils import predict  # noqa: E402  （第 9 讲：model.eval() + no_grad 的批量预测）

sys.path.insert(0, HERE)
from shapes_world import SOURCE_CLASSES, channel_stats, make_shapes, normalize  # noqa: E402


# =============================================================================
# 1. 训练与评估
# =============================================================================
def train_steps(model, X, y, steps, lr=1e-3, bs=32, seed=0, optimizer="adam", augment=None, params=None,
                frozen=(), momentum=0.9, record=False):
    """
    用 mini-batch 训练 model 固定的 steps 步。

    【思路】
        · 每个 epoch 开始时用自己的 torch.Generator 打乱顺序（可复现），按 bs 切 batch；
          最后一个不满的 batch 丢掉（N > bs 时），免得出现只有 1 张图的 batch（BN 在训练模式下没法对 1 个样本求方差）；
        · augment 不是 None 时，每个 batch 先过 augment(xb, gen) —— 每一步看到的都是"新"图（讲义 §8）；
        · params 不是 None 时，只训练这些参数（其余参数不交给优化器，也就不会被更新）；
        · frozen 里的子模块在 model.train() 之后立刻切回 eval()，它们的 BN 用预训练时的统计量、不再更新。
    【参数】
        model     : nn.Module
        X, y      : numpy 数组或张量，X 形状 (N, C, H, W)（已经标准化），y 形状 (N,) 整数类别
        steps     : 梯度更新的步数
        lr        : 学习率；optimizer："adam" 或 "sgd"（带动量 momentum）
        bs        : batch 大小（N ≤ bs 时每一步都用全部 N 张）
        seed      : 打乱顺序和数据增强用的种子
        augment   : None 或函数 f(xb, gen) → xb'
        params    : None（训练全部参数）或参数的可迭代对象，也可以是 [{"params": ..., "lr": ...}, ...] 这种分组
        frozen    : 要保持 eval 模式的子模块列表
        record    : True 时记录每一步的训练损失
    【返回】dict：{"losses": 每步损失的 list（record=False 时只记最后 50 步）, "diverged": 是否出现 inf / nan}
    【边界】损失变成 inf / nan 就立刻停止（第 1 讲 §6.3 的教训：带着 nan 训练下去，结果会"安静地"出错）。
    【PyTorch 小知识】
        torch.randperm(N, generator=g)：0..N−1 的随机排列；X[idx] 用下标张量一次取出一批样本
        torch.optim.Adam(params, lr)  ：params 可以是参数列表，也可以是"参数组"列表（每组自己的学习率）
    """
    X = torch.as_tensor(X)
    y = torch.as_tensor(y)
    gen = torch.Generator().manual_seed(seed)
    if params is None:
        params = [p for p in model.parameters() if p.requires_grad]
    if optimizer == "adam":
        opt = torch.optim.Adam(params, lr=lr)
    else:
        opt = torch.optim.SGD(params, lr=lr, momentum=momentum)
    loss_fn = nn.CrossEntropyLoss()
    N = len(X)
    losses, step = [], 0
    while step < steps:
        model.train()
        for module in frozen:
            module.eval()
        perm = torch.randperm(N, generator=gen)
        last = N - bs + 1 if N > bs else 1                    # 丢掉最后一个不满的 batch
        for start in range(0, last, bs):
            idx = perm[start:start + bs]
            xb, yb = X[idx], y[idx]
            if augment is not None:
                xb = augment(xb, gen)
            logits = model(xb)
            loss = loss_fn(logits, yb)
            opt.zero_grad()
            loss.backward()
            opt.step()
            value = loss.item()
            if not np.isfinite(value):
                return {"losses": losses + [value], "diverged": True}
            if record or step >= steps - 50:
                losses.append(value)
            step += 1
            if step >= steps:
                break
    return {"losses": losses, "diverged": False}


def accuracy(model, X, y):
    """
    分类准确率：argmax(logits) == y 的比例。

    【参数】model；X：(N, ...) numpy 数组或张量；y：(N,) 整数
    【返回】float
    【思路】用第 9 讲的 torch_utils.predict：它会 model.eval() + no_grad，分批前向。
    """
    logits = predict(model, X)
    return float((logits.argmax(dim=1).numpy() == np.asarray(y)).mean())


def features(module, X):
    """
    把 X 送进 module（通常是冻结的网络前半段），返回输出的 numpy 数组。

    【DL 联系】Ng 的建议：数据很少、要冻结整个底座时，可以**预先算好**每张训练图的特征存起来，
            然后只在特征上训练一个小的分类器 —— 每个 epoch 不用再把图像过一遍底座，快得多（讲义 §7.3）。
            代价是没法做"实时"数据增强（每次增强后的图都不一样，特征要重新算）。
    【返回】numpy 数组，形状 (N, ...)
    """
    return predict(module, X).numpy()


# =============================================================================
# 2. 迁移学习用的小 CNN
# =============================================================================
class TransferCNN(nn.Module):
    """
    迁移学习实验用的小 CNN（输入 3×24×24）：

        block1 ：conv 3×3, 16 → BN → ReLU → 2×2 最大池化                 → 16×12×12
        block2 ：(conv 3×3 → BN → ReLU) × 2，32 通道 → 2×2 最大池化       → 32×6×6
        block3 ：(conv 3×3 → BN → ReLU) × 2，64 通道 → 全局平均池化 → 拉平  → 64
        head   ：全连接 64 → n_classes（输出 logits）

    【思路】把网络分成"底座"（base：block1–3，提特征）和"头"（head：分类器）。迁移学习时：
            · 换一个新的 head（目标任务的类别数不同）；
            · 决定冻结到哪一块：只训练 head / 再加上 block3 / 全部都训练（讲义 §7.3）。
    【参数】n_classes：类别数；widths：三个块的通道数
    【PyTorch 小知识】nn.Module 里把子模块赋给 self.xxx，它就会被登记；model.block3.parameters() 只取这一块的参数。
    """

    def __init__(self, n_classes, widths=(16, 32, 64)):
        super().__init__()
        a, b, c = widths

        def cbr(ci, co):
            return [nn.Conv2d(ci, co, 3, padding=1, bias=False), nn.BatchNorm2d(co), nn.ReLU()]

        self.block1 = nn.Sequential(*cbr(3, a), nn.MaxPool2d(2))
        self.block2 = nn.Sequential(*cbr(a, b), *cbr(b, b), nn.MaxPool2d(2))
        self.block3 = nn.Sequential(*cbr(b, c), *cbr(c, c), nn.AdaptiveAvgPool2d(1), nn.Flatten())
        self.head = nn.Linear(c, n_classes)

    def base(self, x):
        """底座：图像 → 64 维特征。"""
        return self.block3(self.block2(self.block1(x)))

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


class Sub(nn.Module):
    """
    把若干个子模块按顺序串起来当作一个新模块（不复制参数，和原模型共享）。

    【用法】Sub(model.block1, model.block2) 就是"原模型的前两块"；features(Sub(...), X) 取出中间特征。
    【PyTorch 小知识】nn.ModuleList 像 list 一样存子模块，但里面的模块会被 PyTorch 登记。
    """

    def __init__(self, *modules):
        super().__init__()
        self.parts = nn.ModuleList(modules)

    def forward(self, x):
        for part in self.parts:
            x = part(x)
        return x


def pretrain_source(seed=0, n_per_class=600, epochs=4, size=24, lr=1e-3, bs=64):
    """
    在源任务（16 种形状，每类 n_per_class 张）上预训练一个 TransferCNN —— 本讲的"ImageNet 预训练模型"。

    【思路】源数据集比任何一个目标数据集都大得多（默认 9600 张 vs 目标任务最多 1500 张），类别也多。
            训练完之后，block1–3 学到的是"形状世界通用的特征"（边缘、角、弧线、斑点……），
            这正是我们希望迁移到新任务上的东西。
    【参数】seed：数据和初始化的种子；n_per_class、epochs、size、lr、bs：见上
    【返回】dict：
        "model"     ：训练好的 TransferCNN（eval 模式）
        "mean"、"std"：源训练集的通道均值 / 标准差 —— 目标任务也必须用**这一组**做标准化
        "train_acc"、"test_acc"：源任务上的准确率（测试集每类 200 张，独立生成）
    【边界】在单线程 CPU 上约 12 秒。固定种子，重复运行结果逐位相同。
    【DL 联系】真实世界里你不会自己做这一步，而是下载别人在 ImageNet 上训练好的权重（讲义 §7.2）。
    """
    X, y = make_shapes(n_per_class, SOURCE_CLASSES, np.random.default_rng(1000 + seed), size=size)
    Xt, yt = make_shapes(200, SOURCE_CLASSES, np.random.default_rng(2000 + seed), size=size)
    mean, std = channel_stats(X)
    Xn, Xtn = normalize(X, mean, std), normalize(Xt, mean, std)
    torch.manual_seed(seed)
    model = TransferCNN(len(SOURCE_CLASSES))
    steps = epochs * (len(X) // bs)
    info = train_steps(model, Xn, y, steps, lr=lr, bs=bs, seed=seed)
    model.eval()
    return {"model": model, "mean": mean, "std": std, "diverged": info["diverged"],
            "train_acc": accuracy(model, Xn[:2000], y[:2000]), "test_acc": accuracy(model, Xtn, yt),
            "n_train": len(X), "steps": steps}


# =============================================================================
# 3. 数据增强（都在一个 batch 上做；输入是标准化之后的张量 (B, 3, H, W)）
# =============================================================================
def flip_lr(xb, gen):
    """
    每张图以 1/2 的概率左右翻转（mirroring）。

    【PyTorch 小知识】xb.flip(-1) 把最后一维（宽）倒过来；torch.where(条件, a, b) 逐元素地"条件真取 a、否则取 b"，
            条件的形状 (B, 1, 1, 1) 会广播到 (B, 3, H, W)。
    """
    mask = torch.rand(len(xb), generator=gen) < 0.5
    return torch.where(mask[:, None, None, None], xb.flip(-1), xb)


def translate(xb, gen, pad=3):
    """
    随机平移 = 先在四周补 pad 圈（复制边缘像素），再随机裁出原来大小的一块（random crop）。

    【思路】这就是常说的 "padding + random crop"：裁剪框的左上角在 [0, 2·pad] 里随机取，相当于把图平移 −pad..pad 像素。
    【参数】pad：最大平移像素数
    【PyTorch 小知识】F.pad(x, (左, 右, 上, 下), mode="replicate") 在最后两维四周补边，"replicate" 表示复制边缘值。
    """
    B, _, H, W = xb.shape
    big = F.pad(xb, (pad, pad, pad, pad), mode="replicate")
    offs = torch.randint(0, 2 * pad + 1, (B, 2), generator=gen)
    return torch.stack([big[i, :, int(oy):int(oy) + H, int(ox):int(ox) + W] for i, (oy, ox) in enumerate(offs)])


def color_jitter(xb, gen, mean, std, strength=0.3):
    """
    颜色扰动：每张图的每个通道乘一个随机增益、加一个随机偏移（在原始像素尺度上），模拟光照和白平衡的变化。

    【思路】原始像素 x = x̂ · σ + μ（x̂ 是标准化后的值）；做 x' = g · x + b，g ∈ [1 − s, 1 + s]，b ∈ [−s/2, s/2]；
            再标准化回去。增益和偏移对 R、G、B 各抽一个，所以颜色的"色调"也会变。
    【参数】mean、std：(3,) 标准化用的通道均值 / 标准差；strength：s
    """
    B = len(xb)
    m = torch.as_tensor(mean)[None, :, None, None]
    s = torch.as_tensor(std)[None, :, None, None]
    g = 1 + strength * (2 * torch.rand(B, 3, 1, 1, generator=gen) - 1)
    b = 0.5 * strength * (2 * torch.rand(B, 3, 1, 1, generator=gen) - 1)
    raw = xb * s + m
    return ((g * raw + b) - m) / s


def pca_basis(X_raw):
    """
    AlexNet 的 "PCA 颜色增强" 要用的主成分：所有训练图所有像素的 RGB 值的 3×3 协方差矩阵的特征分解。

    【参数】X_raw：(N, 3, H, W) 未标准化的图
    【返回】(eigval, eigvec)：(3,) 特征值（从大到小）、(3, 3) 特征向量（每一列一个）
    【numpy 小知识】np.linalg.eigh 用于对称矩阵，特征值从小到大排列，所以要倒过来；np.cov 默认每一行是一个变量。
    """
    pixels = X_raw.transpose(1, 0, 2, 3).reshape(3, -1)          # (3, N·H·W)：每一行是一个颜色通道
    eigval, eigvec = np.linalg.eigh(np.cov(pixels))
    order = np.argsort(eigval)[::-1]
    return eigval[order].astype(np.float32), eigvec[:, order].astype(np.float32)


def pca_color(xb, gen, mean, std, eigval, eigvec, sigma=0.1):
    """
    AlexNet 的 PCA 颜色增强（"fancy PCA"）：给每张图的**所有像素**加同一个 RGB 偏移
        Δ = [p1, p2, p3] · [α1·√λ1, α2·√λ2, α3·√λ3]ᵀ，α_i ~ N(0, σ²)。

    【思路】p_i、λ_i 是 pca_basis 算出的主成分和特征值。偏移沿着"训练集里颜色本来就会变化的方向"走，
            变化最大的方向（通常是整体亮度）走得最远。Krizhevsky et al. 说这一招把 top-1 错误率降了 1 个百分点以上。
    【边界】论文的公式写的是 α_i·λ_i；我们按像素在 [0, 1] 尺度上、用标准差 √λ_i（σ 就是"沿主成分走几个标准差"），
            这也是很多开源实现的做法。σ = 0.1 是论文的值。
    """
    B = len(xb)
    alpha = sigma * torch.randn(B, 3, generator=gen)
    delta = (alpha * torch.as_tensor(np.sqrt(eigval))) @ torch.as_tensor(eigvec).T     # (B, 3)
    return xb + (delta / torch.as_tensor(std))[:, :, None, None]


def compose(*fns):
    """
    把几个增强函数串起来：compose(f, g)(xb, gen) = g(f(xb, gen), gen)。

    【Python 小知识】*fns 把任意多个参数收集成一个元组；返回的是一个"闭包"函数，它记住了 fns。
    """
    def run(xb, gen):
        for fn in fns:
            xb = fn(xb, gen)
        return xb
    return run


if __name__ == "__main__":
    torch.set_num_threads(1)
    a = pretrain_source()
    b = pretrain_source()
    same = all(torch.equal(p, q) for p, q in zip(a["model"].state_dict().values(), b["model"].state_dict().values()))
    print(f"源任务预训练：{a['n_train']} 张、{a['steps']} 步；训练准确率 {a['train_acc']:.3f}，测试 {a['test_acc']:.3f}；"
          f"再训练一次参数逐位相同：{same}")
    gen = torch.Generator().manual_seed(0)
    xb = torch.randn(8, 3, 24, 24)
    for name, fn in [("flip_lr", flip_lr), ("translate", translate),
                     ("color_jitter", lambda x, g: color_jitter(x, g, a["mean"], a["std"]))]:
        out = fn(xb, gen)
        print(f"    {name:<13} 形状 {tuple(out.shape)}，和原图的平均绝对差 {float((out - xb).abs().mean()):.3f}")