"""
=============================================================================
CS230 · 第 12 讲工具模块：小 CNN 的 PyTorch 模型、训练封装、数据与字符画
=============================================================================

对应讲义：lectures/lecture12_cnn_foundations.md §8、§9

-----------------------------------------------------------------------------
【这个模块里有什么】
-----------------------------------------------------------------------------
    数据
        cross_split            第 2 讲的"橙色十字"数据（16×16×3），整理成 PyTorch 的 (m, 3, 16, 16) 形状
        make_crosses_at        同一种图，但图案的位置由我们指定（用来做"只在左边训练、在右边测试"的极端实验）
        hand_sign_tensors      第 9 讲的合成手势数据（16×16×1），整理成 (m, 1, 16, 16)
    模型
        make_mlp               先拉平、再接第 9 讲的 MLP（全连接网络）
        SmallCNN               CONV → ReLU → MAXPOOL → CONV → ReLU → MAXPOOL → 输出头（"fc" 拉平接全连接 / "gmp" 全局最大池化）
    训练与评估
        train_classifier       固定种子 → 建模型 → 用第 9 讲的 fit() 训练（Adam）
        accuracy               分类准确率
        describe_layers        逐层列出输出形状、激活值个数、参数个数（讲义 §8.1 的那种表）
    字符画
        unsigned_heatmap       非负的图（比如 ReLU 之后的特征图）画成字符
        side_by_side           把几块字符画并排打印

依赖前几讲的模块（只读导入）：
    第 1 讲 console.py            pad、print_header
    第 2 讲 image_data.py         橙色十字数据：load_dataset、draw_background、plus_mask、square_mask、COLORS……；signed_heatmap
    第 2 讲 text_plot.py          matplotlib 的统一设置（get_pyplot、style_axes、figure_path、SERIES）
    第 9 讲 torch_utils.py        MLP、make_loader、fit、predict、count_parameters、setup_torch
    第 9 讲 multiclass_data.py    hand_sign_splits、make_hand_signs、CLASS_NAMES

【形状约定】PyTorch 的习惯：图像是 (m, C, H, W)（通道在前）；标签 y 是 (m,) 的整数类别。
            二分类也用 2 个输出 + nn.CrossEntropyLoss（第 9 讲 §3.4：C = 2 的 softmax 就是逻辑回归），
            这样所有模型共用第 9 讲的训练循环。代价是输出层的参数多记了一份（2 个输出而不是 1 个），对所有模型一视同仁。

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

import os
import sys

import numpy as np
import torch
from torch import nn

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

from console import pad, print_header  # noqa: E402,F401  （第 1 讲：对齐表格、分隔标题；脚本从这里导入）
import image_data  # noqa: E402            （第 2 讲：橙色十字数据）
from image_data import signed_heatmap  # noqa: E402,F401
from multiclass_data import CLASS_NAMES, hand_sign_splits  # noqa: E402,F401  （第 9 讲：合成手势）
from text_plot import SERIES, figure_path, get_pyplot, style_axes  # noqa: E402,F401  （第 2 讲：画图设置）
from torch_utils import MLP, count_parameters, fit, make_loader, predict, setup_torch  # noqa: E402,F401  （第 9 讲）


# =============================================================================
# 1. 数据
# =============================================================================
def to_nchw(X_orig):
    """
    (m, H, W, C) 的 uint8 图像 → (m, C, H, W) 的 float32，像素除以 255 缩放到 [0, 1]。

    【思路】第 2 讲的图像是"通道在最后"（和 Coursera 一样），PyTorch 的卷积层要"通道在前"。
            用 transpose 换轴（不能用 reshape，讲义 §6.4）。
    【参数】X_orig：(m, H, W, C)
    【返回】(m, C, H, W) float32
    """
    return np.ascontiguousarray(X_orig.transpose(0, 3, 1, 2)).astype(np.float32) / 255.0


def cross_split(jitter, seed=0, m_train=1000, m_test=500):
    """
    第 2 讲的"橙色十字 vs 其他"数据：和第 2 讲 §9.3 用的是同一份（同样的种子、同样的大小）。

    【参数】jitter：位置抖动幅度 0–4；seed、m_train、m_test：传给 image_data.load_dataset
    【返回】dict：
        "train" : (X, y)，X (m_train, 3, 16, 16) float32，y (m_train,) int64
        "test"  : (X, y)
        "test_kind" : (m_test,) 字符串，每张测试图的种类（plus / orange_square / other_plus / empty），只用于分析
    """
    d = image_data.load_dataset(jitter=jitter, m_train=m_train, m_test=m_test, seed=seed)
    return {"train": (to_nchw(d["train_x_orig"]), d["train_y"][0].astype(np.int64)),
            "test": (to_nchw(d["test_x_orig"]), d["test_y"][0].astype(np.int64)),
            "test_kind": d["test_kind"]}


def _draw_shape_at(img, mask, color_name, rng, top, left):
    """
    把图案画在**指定的**位置 (top, left)。颜色的随机扰动和第 2 讲 image_data.draw_shape 完全一样，
    只是位置不再是"中间 ± 抖动"，而是由调用者决定。

    【参数】img：(16, 16, 3) 浮点数组（原地修改）；mask：(7, 7) 布尔蒙版；color_name：颜色名；
            rng：随机数生成器；top、left：图案左上角，0 ≤ top, left ≤ 9（16 − 7 = 9）
    """
    size = image_data.SHAPE_SIZE
    color = np.array(image_data.COLORS[color_name], dtype=float) + rng.normal(0, 12, 3)
    sub = img[top:top + size, left:left + size]
    sub[mask] = color + rng.normal(0, 20, (int(mask.sum()), 3))


def make_crosses_at(m, rng, lefts, tops):
    """
    生成 m 张"橙色十字 vs 其他"图，图案左上角的列从 lefts 里随机选、行从 tops 里随机选。

    【思路】和第 2 讲 image_data.make_images 一样：一半是正类（橙色十字），另一半在三种干扰项
            （橙色方块 / 其他颜色的十字 / 只有背景）里均匀选；背景、颜色扰动、噪声都用第 2 讲的函数。
            唯一的区别是位置：例如 lefts = [0, 1, 2] 时，图案只出现在图的**左边**（占第 0–8 列）。
    【参数】m：张数；rng：np.random.Generator；lefts、tops：可选的列 / 行起点（整数列表，每个在 0–9 之间）
    【返回】(X, y)：X (m, 3, 16, 16) float32 ∈ [0, 1]，y (m,) int64
    【边界】起点超过 9 时图案会超出 16×16 的画布，这里直接报错。
    【DL 联系】讲义 §9.5 的极端测试：训练时图案只在左边，测试时只在右边 —— 哪种结构能"举一反三"？
    """
    if max(lefts) > 9 or max(tops) > 9 or min(lefts) < 0 or min(tops) < 0:
        raise ValueError("图案左上角必须在 0–9 之间")
    X = np.zeros((m, 16, 16, 3), dtype=np.uint8)
    y = np.zeros(m, dtype=np.int64)
    for i in range(m):
        img = image_data.draw_background(rng)
        top = int(tops[int(rng.integers(len(tops)))])
        left = int(lefts[int(rng.integers(len(lefts)))])
        if rng.random() < 0.5:
            _draw_shape_at(img, image_data.plus_mask(), "orange", rng, top, left)
            y[i] = 1
        else:
            kind = image_data.KINDS[1 + int(rng.integers(3))]
            if kind == "orange_square":
                _draw_shape_at(img, image_data.square_mask(), "orange", rng, top, left)
            elif kind == "other_plus":
                color = image_data.OTHER_COLORS[int(rng.integers(len(image_data.OTHER_COLORS)))]
                _draw_shape_at(img, image_data.plus_mask(), color, rng, top, left)
        X[i] = np.round(np.clip(img, 0, 255)).astype(np.uint8)
    return to_nchw(X), y


def hand_sign_tensors(standardize="channel"):
    """
    第 9 讲的合成手势数据（训练 1080 / dev 500 / test 1000），整理成 (m, 1, 16, 16)，并做输入标准化。

    【思路】两种标准化（都只用训练集算统计量）：
            · "pixel"   ：每个像素一个均值、一个标准差（第 9 讲代码 05 对 MLP 的做法）；
            · "channel" ：整个通道共用一个均值、一个标准差（CNN 的常见做法）。
            为什么 CNN 更常用后者？卷积的想法是"同一个检测器在每个位置做同样的事"；如果每个像素先被
            减去不同的均值、除以不同的标准差，同一个图案出现在不同位置时，进入网络的数字就不一样了。
            对全连接网络这无所谓（它本来就给每个位置单独的权重）。
    【参数】standardize："pixel" 或 "channel"
    【返回】dict：{"train": (X, y), "dev": (X, y), "test": (X, y)}，X (m, 1, 16, 16) float32，y (m,) int64
    """
    raw = hand_sign_splits()
    Xtr = raw["train"][0]
    if standardize == "pixel":
        mu, sd = Xtr.mean(axis=0), Xtr.std(axis=0) + 1e-8          # 形状 (16, 16)
    elif standardize == "channel":
        mu, sd = Xtr.mean(), Xtr.std()                              # 两个数
    else:
        raise ValueError(f"unknown standardize: {standardize}")
    return {name: (((X - mu) / sd)[:, None, :, :].astype(np.float32), y) for name, (X, y) in raw.items()}


# =============================================================================
# 2. 模型
# =============================================================================
def make_mlp(layer_dims):
    """
    全连接网络：先把 (m, C, H, W) 的图拉平成 (m, C·H·W)，再接第 9 讲 torch_utils 的 MLP。

    【参数】layer_dims：[n_x, n[1], ..., n[L]]，n_x 必须等于 C·H·W
    【返回】nn.Sequential(nn.Flatten(), MLP(layer_dims))
    【PyTorch 小知识】nn.Flatten() 默认把第 1 维之后的所有维拉平：(m, 3, 16, 16) → (m, 768)。
            它和第 1 讲的 reshape(m, -1) 一样，只是顺序是 (C, H, W) 而不是 (H, W, C) —— 对全连接网络没有区别。
    """
    return nn.Sequential(nn.Flatten(), MLP(layer_dims))


class SmallCNN(nn.Module):
    """
    一个最简单的卷积网络：CONV → ReLU → MAXPOOL → CONV → ReLU → MAXPOOL → 输出头。

    【结构】（输入 in_ch × img × img，例如 3 × 16 × 16）
            CONV1 : c1 个 3×3 滤波器，stride 1，padding 1（"same"：输出还是 img × img）
            MAXPOOL : 2×2，stride 2 → img/2 × img/2
            CONV2 : c2 个 3×3×c1 的滤波器，same
            MAXPOOL : 2×2，stride 2 → img/4 × img/4
            输出头 head：
                "fc"  ：拉平成 c2·(img/4)² 个数，接一个全连接层 → n_classes 个 logits（可选再加一层隐藏层 fc_hidden）
                "gmp" ：全局最大池化（global max pooling）：每个通道只保留整张图上的最大值 → c2 个数 → 全连接 → logits
    【思路】两个卷积层负责"在任何位置找局部图案"（平移等变），两个池化层把图逐步缩小（讲义 §7、§9.3）。
            "fc" 头是 Ng 讲的标准结构（CONV-POOL-CONV-POOL-FC，讲义 §8）；它保留了"图案在哪个格子里"的信息，
            所以整个网络并不是完全平移不变的。"gmp" 头把位置信息彻底丢掉（"有没有"而不是"在哪里"），
            讲义 §9.5 的极端测试会比较两者。
    【参数】in_ch：输入通道数；n_classes：类别数；c1、c2：两个卷积层的滤波器个数；img：输入边长（要能被 4 整除）；
            head："fc" 或 "gmp"；fc_hidden：fc 头里隐藏层的单元数（0 表示没有隐藏层）
    【PyTorch 小知识】
        nn.Conv2d(in, out, kernel_size=3, padding=1)：weight 形状 (out, in, 3, 3)，bias 形状 (out,)。
        nn.MaxPool2d(2)：窗口 2×2，stride 默认等于窗口大小。
        nn.AdaptiveMaxPool2d(1)：不管输入多大，都把每个通道池化成 1×1（就是取整张图的最大值）。
    """

    def __init__(self, in_ch, n_classes, c1=8, c2=16, img=16, head="fc", fc_hidden=0):
        super().__init__()
        if img % 4 != 0:
            raise ValueError("img 要能被 4 整除（两次 2×2 池化）")
        self.features = nn.Sequential(
            nn.Conv2d(in_ch, c1, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(c1, c2, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        )
        if head == "fc":
            n_flat = c2 * (img // 4) ** 2
            if fc_hidden > 0:
                self.head = nn.Sequential(nn.Flatten(), nn.Linear(n_flat, fc_hidden), nn.ReLU(), nn.Linear(fc_hidden, n_classes))
            else:
                self.head = nn.Sequential(nn.Flatten(), nn.Linear(n_flat, n_classes))
        elif head == "gmp":
            self.head = nn.Sequential(nn.AdaptiveMaxPool2d(1), nn.Flatten(), nn.Linear(c2, n_classes))
        else:
            raise ValueError(f"unknown head: {head}")

    def forward(self, x):
        """x：(m, in_ch, img, img) → logits：(m, n_classes)。"""
        return self.head(self.features(x))


# =============================================================================
# 3. 训练与评估
# =============================================================================
def train_classifier(build_fn, X, y, seed, epochs, lr=1e-3, batch_size=32):
    """
    固定种子 → 建模型 → 用 Adam 训练 epochs 个 epoch（第 9 讲的 fit()，不用 dev 集、不早停）。

    【思路】种子要在**建模型之前**设（层的初始化用 PyTorch 的全局随机数，第 9 讲 §4.6）；
            DataLoader 打乱顺序用自己的 torch.Generator（同一个 seed）。这样同一个 (seed, 配置) 每次结果逐位相同。
    【参数】
        build_fn   : 一个不带参数的函数，调用它返回一个新模型（例如 lambda: SmallCNN(3, 2)）
        X, y       : 训练数据，X (m, ...) float32，y (m,) int64
        seed       : 种子
        epochs     : epoch 数
        lr         : Adam 的学习率
        batch_size : mini-batch 大小
    【返回】(model, history, info)：训练好的模型；fit() 的历史和信息（info["diverged"] 表示是否发散）
    【Python 小知识】把"怎样建模型"当成一个函数传进来（build_fn），这样同一个训练函数可以训练任何模型。
            lambda: SmallCNN(3, 2) 是一个"没有参数、返回 SmallCNN(3, 2)"的小函数。
    """
    torch.manual_seed(seed)
    model = build_fn()
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    hist, info = fit(model, opt, make_loader(X, y, batch_size, shuffle=True, seed=seed), None, epochs=epochs)
    return model, hist, info


def accuracy(model, X, y):
    """分类准确率：argmax(logits) == y 的比例。X (m, ...)，y (m,)。返回 Python float。"""
    pred = predict(model, X).argmax(dim=1).numpy()
    return float(np.mean(pred == y))


def _describe(module):
    """给一层起一个人能读懂的名字（讲义的写法：CONV / POOL / FC）。"""
    if isinstance(module, nn.Conv2d):
        k, s, p = module.kernel_size[0], module.stride[0], module.padding[0]
        return f"CONV {k}×{k}，{module.out_channels} 个滤波器，s={s}，p={p}"
    if isinstance(module, nn.MaxPool2d):
        k = module.kernel_size if isinstance(module.kernel_size, int) else module.kernel_size[0]
        s = module.stride if isinstance(module.stride, int) else module.stride[0]
        return f"MAXPOOL f={k}，s={s}"
    if isinstance(module, nn.AvgPool2d):
        return f"AVGPOOL f={module.kernel_size}，s={module.stride}"
    if isinstance(module, nn.AdaptiveMaxPool2d):
        return "全局 MAXPOOL"
    if isinstance(module, nn.Linear):
        return f"FC {module.in_features} → {module.out_features}"
    return type(module).__name__                    # ReLU、Flatten……


def describe_layers(model, input_shape, skip=("ReLU", "Flatten")):
    """
    把一个样本送进网络，逐层记录输出的形状、激活值个数和这一层的参数个数。

    【思路】给每个"叶子"层（没有子层的层）挂一个 forward hook：前向传播经过这一层时，hook 被调用，
            我们就把这一层的输出形状记下来。前向传播一次就得到整张表。
    【参数】model：nn.Module；input_shape：一个样本的形状，例如 (3, 32, 32)；skip：不单独列出的层的类型名
    【返回】list[dict]，每个 dict 有 "layer"（描述）、"shape"（不含 batch 维的形状）、"size"（激活值个数）、"params"
            第一行是输入本身（参数 0）。
    【PyTorch 小知识】
        module.register_forward_hook(fn)：以后每次 module 做前向，都会调用 fn(module, 输入, 输出)。
            它返回一个 handle，用完要 handle.remove()，否则 hook 会一直挂着。
        model.modules()：按定义的顺序遍历所有层（包括嵌套在 nn.Sequential 里的）。
    """
    rows = [{"layer": "输入", "shape": tuple(input_shape), "size": int(np.prod(input_shape)), "params": 0}]
    handles = []

    def hook(module, inputs, output):
        if type(module).__name__ in skip:
            return
        shape = tuple(output.shape[1:])
        rows.append({"layer": _describe(module), "shape": shape, "size": int(np.prod(shape)),
                     "params": sum(p.numel() for p in module.parameters())})

    for module in model.modules():
        if len(list(module.children())) == 0:           # 只给叶子层挂 hook
            handles.append(module.register_forward_hook(hook))
    model.eval()
    with torch.no_grad():
        model(torch.zeros((1,) + tuple(input_shape)))
    for h in handles:
        h.remove()
    return rows


def print_layer_table(rows, shape_fmt="hwc"):
    """
    打印 describe_layers 的结果。shape_fmt = "hwc" 时把 (C, H, W) 写成 Ng 的 H×W×C；一维的写成 (n,)。

    【返回】参数总数
    """
    print(pad("层", 34) + pad("输出形状（高×宽×通道）", 26, ">") + pad("激活值个数", 12, ">") + pad("参数个数", 12, ">"))
    total = 0
    for r in rows:
        s = r["shape"]
        if len(s) == 3 and shape_fmt == "hwc":
            shape_txt = f"{s[1]}×{s[2]}×{s[0]}"
        else:
            shape_txt = f"({s[0]},)" if len(s) == 1 else str(s)
        print(pad(r["layer"], 34) + pad(shape_txt, 26, ">") + f"{r['size']:>12,}{r['params']:>12,}")
        total += r["params"]
    print(pad("合计", 34) + " " * 38 + f"{total:>12,}")
    return total


# =============================================================================
# 4. 字符画
# =============================================================================
def unsigned_heatmap(M, scale=None, chars=" .:-=+*#%@"):
    """
    把一个非负的二维数组（比如 ReLU 之后的特征图）画成字符：越大越"浓"，每个值占 2 个字符（看起来更接近正方形）。

    【参数】M：(h, w)，值 ≥ 0；scale：归一化用的最大值（默认是 M 自己的最大值；几张图要互相比较时传同一个 scale）
    【返回】list[str]，h 行
    【边界】M 全为 0 时 scale 兜底为 1e−12，所有格子都画成最淡的字符。
    """
    if scale is None:
        scale = float(np.max(M))
    scale = max(scale, 1e-12)
    n = len(chars)
    idx = np.clip((np.asarray(M) / scale * (n - 1)).round().astype(int), 0, n - 1)
    return ["".join(chars[v] * 2 for v in row) for row in idx]


def side_by_side(blocks, titles, gap=4):
    """
    把几块字符画（每块是 list[str]）并排打印，每块上面有一个标题。

    【参数】blocks：list[list[str]]；titles：list[str]，和 blocks 一一对应；gap：两块之间的空格数
    【返回】list[str]，可以直接逐行 print
    【Python 小知识】zip(*blocks) 把"每块的第 r 行"凑成一组：先取所有块的第 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