"""
=============================================================================
第 4 讲 · 代码 4：多层感知机与手写 backprop（MLP, backpropagation, XOR, initialization）
=============================================================================

对应讲义：lectures/lecture04_trees_ensembles_nn.md 第 5~6 节

-----------------------------------------------------------------------------
【这个文件要回答的问题】
-----------------------------------------------------------------------------
第 1 讲 §4.7 说：神经网络 = 学出来的特征 φ(x) + 一个 GLM 输出头，backprop 从 logits 上的 (ŷ − y) 开始往回传。
这里把这句话变成能运行、能检查的代码。

    演示 1：梯度检查（gradient check）。手写的 backward 和中心差分逐个比较；logits 上的梯度确实是 (P − Y)/n；
            ReLU 的“折角”会让梯度检查在个别位置失败（不是 bug）；可选：和 PyTorch autograd 对照。
    演示 2：XOR —— 感知机的滑铁卢（第 1 讲 §3.6）。感知机永远停不下来，逻辑回归卡在 θ = 0；
            一个隐藏层就能解决，但隐藏单元太少时经常训练失败（20 个随机种子的成功率）。打印隐藏层学到的特征。
    演示 3：初始化。(a) 全 0 / 全相同的初始化：隐藏单元永远一模一样（对称性）；
            (b) 深层网络里权重的尺度：太小 → 信号和梯度消失，太大 → 饱和 / 爆炸；gain = 1（tanh）、√2（ReLU）刚好。
    演示 4：backprop 的代价。数一数乘加次数：反向 ≈ 2 × 前向，和参数个数无关；数值梯度要 2P 次前向。
            向量化：一次矩阵乘法 = 把每个样本的梯度加起来。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 04_mlp_backprop.py          # 约 8 秒（主要是演示 2 的 200 次 XOR 训练）

依赖本目录的 neural_net.py、nonlinear_data.py，以及第 1 讲的 linear_models.py。
PyTorch 是可选的：装了就多一项对照，没装就跳过（不影响其他输出）。
"""

import numpy as np

from neural_net import (backward, count_macs, forward, hidden_features, init_params, loss_and_grad, n_params,
                        numerical_grad, predict, relative_error, train)
from nonlinear_data import make_circle, xor_points
from linear_models import accuracy, add_intercept, fit_logistic_newton, logistic_grad, perceptron, print_header


# =============================================================================
# 演示 1：梯度检查
# =============================================================================
def torch_check(params, X, y, act, l2):
    """
    用 PyTorch autograd 算同一个网络、同一个损失的梯度，返回和手写 backprop 的最大绝对差；没装 torch 时返回 None。

    【思路】把 numpy 参数原样拷成 float64 的 torch 张量（requires_grad=True），用 torch 的算子重写前向，
        调用 loss.backward()，torch 自动沿计算图反向求导（reverse-mode autodiff）—— 和我们手写的是同一个算法。
    【边界】torch.set_num_threads(1)：单线程，保证每次运行结果逐位相同。
        import 放在 try 里：没装 torch 时 ImportError，直接返回 None。
    """
    try:
        import torch
    except ImportError:
        return None
    torch.set_num_threads(1)
    T = [torch.tensor(p, dtype=torch.float64, requires_grad=True) for p in params]
    A = torch.tensor(X, dtype=torch.float64)
    f = {"tanh": torch.tanh, "relu": torch.relu, "sigmoid": torch.sigmoid}[act]
    L = len(T) // 2
    for ell in range(L):
        Z = A @ T[2 * ell] + T[2 * ell + 1]
        A = f(Z) if ell < L - 1 else Z
    loss = torch.nn.functional.cross_entropy(A, torch.tensor(y, dtype=torch.long))    # 平均 CE，和我们的一样
    loss = loss + 0.5 * l2 * sum((T[i] ** 2).sum() for i in range(0, len(T), 2))
    loss.backward()
    _, grads, _ = loss_and_grad(params, X, y, act, "softmax", l2)
    return max(float(np.max(np.abs(t.grad.numpy() - g))) for t, g in zip(T, grads))


def demo_gradient_check():
    print_header("演示 1：梯度检查 —— 手写 backprop vs 中心差分（vs PyTorch autograd）")
    rng = np.random.default_rng(41)
    X = rng.normal(size=(8, 3))
    y = rng.integers(0, 3, size=8)
    sizes = [3, 5, 4, 3]

    # (a) logits 上的梯度
    Z = rng.normal(size=(8, 3))

    def ce(Zm):
        zmax = Zm.max(axis=1, keepdims=True)
        lse = zmax[:, 0] + np.log(np.exp(Zm - zmax).sum(axis=1))
        return float(np.mean(lse - Zm[np.arange(8), y]))

    num = numerical_grad([Z], lambda: ce(Z))[0]
    P = np.exp(Z - Z.max(axis=1, keepdims=True))
    P /= P.sum(axis=1, keepdims=True)
    P[np.arange(8), y] -= 1.0
    err_logits = float(np.max(np.abs(num - P / 8)))
    print(f"(a) 平均 cross-entropy 对 logits 的数值梯度 vs (P − Y)/n：最大绝对差 = {err_logits:.1e}\n")

    # (b) 整个网络
    print(f"(b) 网络结构 {sizes}（{n_params(init_params(sizes, rng))} 个参数），8 个样本，L2 = 0.1。")
    print(f"{'activation':<24}{'relative error':>16}{'worst element':>15}{'loss evals (2P)':>17}{'mismatched arrays':>19}")
    res = {}
    for act, kink in [("tanh", False), ("sigmoid", False), ("relu", False), ("relu", True)]:
        params = init_params(sizes, np.random.default_rng(42))
        Xc = X.copy()
        if kink:
            Xc[0] = 0.0                      # 第一个样本是全 0 向量；b₁ = 0 → 第一层的 Z 恰好是 0，正好落在 ReLU 的折角上
        calls = [0]

        def loss_fn():
            calls[0] += 1
            return loss_and_grad(params, Xc, y, act, "softmax", 0.1)[0]

        _, g, _ = loss_and_grad(params, Xc, y, act, "softmax", 0.1)
        g_num = numerical_grad(params, loss_fn)
        rel, worst = relative_error(g, g_num), relative_error(g, g_num, elementwise=True)
        # 哪些参数数组对不上？（|差| > 1e−6 的数组名，用来定位折角的影响）
        names = ["W" + str(i // 2 + 1) if i % 2 == 0 else "b" + str(i // 2 + 1) for i in range(len(params))]
        bad = [nm for nm, a, b in zip(names, g, g_num) if np.max(np.abs(a - b)) > 1e-6]
        label = act + (" (input row = 0)" if kink else "")
        res[label] = (rel, worst, bad)
        n_evals = calls[0]
        print(f"{label:<24}{rel:>16.1e}{worst:>15.1e}{calls[0]:>17}{', '.join(bad) or '-':>19}")
    kink_label = "relu (input row = 0)"

    # (c) PyTorch
    params = init_params(sizes, np.random.default_rng(42))
    diffs = {act: torch_check(params, X, y, act, 0.1) for act in ("tanh", "relu")}
    if diffs["tanh"] is None:
        print("\n(c) 没有安装 PyTorch，跳过 autograd 对照。")
        torch_line = "  · PyTorch 没装，跳过了 (c)。"
    else:
        print(f"\n(c) PyTorch autograd（float64）vs 手写 backprop，最大绝对差：tanh {diffs['tanh']:.1e}，relu {diffs['relu']:.1e}")
        torch_line = (f"  · (c) PyTorch 的 autograd 和手写的 backward 一致到 {max(diffs.values()):.0e}（float64 的舍入误差级别）：\n"
                      "    autograd 做的就是同一件事 —— 前向时记录计算图，反向时对每个节点套用链式法则。")
    print(f"""
观察 Observations：
  · (a) softmax + cross-entropy 对 logits 的梯度就是 (P − Y)/n（差 {err_logits:.0e}）：backprop 的起点是“预测 − 目标”，
    和第 1 讲每个 GLM 的梯度同一个形状。所有的链式法则都从这个干净的信号开始。
  · (b) tanh、sigmoid 和 ReLU（无折角）的整体相对误差都在 1e−9 左右（{min(res[k][0] for k in ('tanh', 'sigmoid', 'relu')):.1e} ~ {max(res[k][0] for k in ('tanh', 'sigmoid', 'relu')):.1e}；
    来自中心差分 O(h²) 的截断误差和舍入误差）：手写的 backward 是对的。逐元素的最大相对误差大一些，因为有些梯度分量本身就接近 0。
  · 把一个输入样本设成全 0 向量后，第一层的 Z 恰好等于 b₁ = 0 —— 正落在 ReLU 的折角上；这个样本在第一层的输出也是 0，
    所以第二层的 Z = b₂ = 0 也落在折角上。ReLU 在 0 处不可导：我们取导数 0，而中心差分算的是 (ReLU(h) − ReLU(−h))/(2h) = 1/2。
    于是整体相对误差变成 {res[kink_label][0]:.0e}，对不上的正好是 {' 和 '.join(res[kink_label][2])} 这两个偏置（扰动 W 不会让 Z 离开 0，因为输入是 0）。
    这不是 bug —— 梯度检查失败时，先看是不是有折角（ReLU、max、abs）。
  · 数值梯度对每个参数要算两次损失：这个 {n_params(params)} 个参数的小网络就要 {n_evals} 次前向。backprop 一次前向 + 一次反向就给出全部梯度。
{torch_line}""")


# =============================================================================
# 演示 2：XOR
# =============================================================================
def demo_xor():
    print_header("演示 2：XOR —— 感知机做不到，逻辑回归做不到，一个隐藏层就够了")
    X, y = xor_points()
    A = add_intercept(X)
    th_p, mistakes = perceptron(A, y, max_epochs=100)
    th_lr, it_lr = fit_logistic_newton(A, y)
    g0 = logistic_grad(np.zeros(3), A, y)
    print("4 个点：(0,0)→0，(0,1)→1，(1,0)→1，(1,1)→0。")
    print(f"  感知机（第 1 讲）：100 个 epoch，每个 epoch 的犯错次数 = {sorted(set(mistakes))} 中的某个数，从未为 0；"
          f"最后的准确率 {accuracy(y, (A @ th_p >= 0).astype(int)):.2f}")
    print(f"  逻辑回归（Newton）：θ = 0 处的梯度 = {np.round(g0, 12).tolist()}，于是 {it_lr} 步就“收敛”到 θ = {th_lr.tolist()}，"
          f"所有预测概率都是 0.5")

    print("\n一个隐藏层的 MLP [2, h, 2]，softmax 输出，Adam（lr = 0.05）训练 1000 步；每个 h 用 20 个随机初始化，")
    print("数一数有几次把 4 个点全部分对：\n")
    print(f"{'h':>3}{'tanh: solved':>14}{'median loss':>13}{'relu: solved':>14}{'median loss':>13}{'relu: all dead':>16}")
    first_solved = None
    table = {}
    for h in [1, 2, 3, 4, 8]:
        row = {}
        for act in ("tanh", "relu"):
            ok, losses = 0, []
            for s in range(20):
                p = train(init_params([2, h, 2], np.random.default_rng(100 + s)), X, y, steps=1000, lr=0.05, act=act)
                solved = accuracy(y, predict(p, X, act)) == 1.0
                ok += solved
                losses.append(loss_and_grad(p, X, y, act)[0])
                if solved and first_solved is None and h == 2 and act == "tanh":
                    first_solved = (s, p)
            # "all dead"：损失 ≈ log 2 = 0.693，说明网络对 4 个点的输出完全相同（ReLU 单元全部恒为 0，只剩偏置）
            row[act] = (ok, float(np.median(losses)), int(np.sum(np.array(losses) > 0.69)))
        table[h] = row
        print(f"{h:>3}{row['tanh'][0]:>11}/20{row['tanh'][1]:>13.4f}{row['relu'][0]:>11}/20{row['relu'][1]:>13.4f}"
              f"{row['relu'][2]:>13}/20")

    s, p = first_solved
    H = hidden_features(p, X)
    logits = forward(p, X)[0]
    print(f"\n第一个解出 XOR 的 h = 2 tanh 网络（种子 {100 + s}）。隐藏层把 4 个点映射成：")
    print(f"{'x':>12}{'y':>4}{'hidden h1':>11}{'hidden h2':>11}{'logit diff z1 - z0':>20}")
    for i in range(4):
        print(f"{str(X[i].astype(int).tolist()):>12}{y[i]:>4}{H[i, 0]:>11.3f}{H[i, 1]:>11.3f}{logits[i, 1] - logits[i, 0]:>20.3f}")
    print(f"""
观察 Observations：
  · 感知机在 XOR 上永远停不下来（每个 epoch 都犯错），最后的准确率只有 50%；逻辑回归更“干脆”：由于对称性，
    θ = 0 处的梯度恰好是 0，Newton 一步都不走。任何线性分类器在 4 个点上最多分对 3 个。
  · h = 1 永远失败（{table[1]['tanh'][0]}/20）：一个隐藏单元只能给出 x 的一个线性投影的单调函数，再接线性输出，还是分不开。
    h = 2 在理论上就够了，但 tanh 只有 {table[2]['tanh'][0]}/20、ReLU 只有 {table[2]['relu'][0]}/20 次成功 —— 损失函数不是凸的，有局部最优 / 平台。
    ReLU 还有一种特有的失败：所有隐藏单元都“死掉”（对 4 个输入的 Z 都 ≤ 0，梯度恒为 0），网络退化成常数，损失 = log 2
    （最后一列；h = 2、4、8 时分别有 {table[2]['relu'][2]}、{table[4]['relu'][2]}、{table[8]['relu'][2]} 次）。
  · 加宽到 h = 8：tanh {table[8]['tanh'][0]}/20、ReLU {table[8]['relu'][0]}/20。**过参数化让优化变容易**：表示 XOR 只需要 2 个单元，但训练时多给一些，
    几乎总能找到解。这是深度学习里反复出现的现象（大模型更好训练），理论上仍不完全清楚。
  · 隐藏层做了什么？看上面的表：两个标签为 0 的点被映射到 (h1, h2) 平面的同一个角落附近，两个标签为 1 的点在另外两个角上 ——
    在隐藏空间里它们**线性可分**了，输出层（一个 softmax 回归 = GLM）只需要画一条直线。“网络 = 学出来的特征 + GLM 头”就是这个意思。""")


# =============================================================================
# 演示 3：初始化
# =============================================================================
def demo_initialization():
    print_header("演示 3：初始化 —— 对称性（symmetry breaking）与尺度（scale）")
    rng = np.random.default_rng(43)
    X, y = make_circle(200, rng)
    Xt, yt = make_circle(2000, rng)
    print("(a) 圆形边界，200 个训练点。网络 [2, 8, 2]，tanh，Adam（lr = 0.02）2000 步。三种初始化：\n")
    print(f"{'init':<22}{'train acc':>11}{'test acc':>10}{'max |W1[:,j] - W1[:,0]|':>25}{'rank of hidden output':>23}")
    res = {}
    for label, kw in [("random N(0, 1/fan_in)", {"mode": "random"}), ("all zeros", {"mode": "zeros"}),
                      ("all W = 0.5", {"mode": "constant", "const": 0.5})]:
        p = train(init_params([2, 8, 2], np.random.default_rng(1), **kw), X, y, steps=2000, lr=0.02)
        spread = float(np.max(np.abs(p[0] - p[0][:, :1])))         # 每个隐藏单元的输入权重和第 0 个单元的差
        rank = int(np.linalg.matrix_rank(hidden_features(p, X), tol=1e-8))
        res[label] = (accuracy(y, predict(p, X)), accuracy(yt, predict(p, Xt)), spread, rank)
        print(f"{label:<22}{res[label][0]:>11.3f}{res[label][1]:>10.3f}{spread:>25.2e}{rank:>23}")
    p1 = train(init_params([2, 1, 2], np.random.default_rng(1)), X, y, steps=2000, lr=0.02)
    acc1 = (accuracy(y, predict(p1, X)), accuracy(yt, predict(p1, Xt)))
    print(f"{'(reference: h = 1)':<22}{acc1[0]:>11.3f}{acc1[1]:>10.3f}")

    print("\n(b) 20 层、每层 100 个单元的网络，输入 500 个 N(0, I) 样本。W ~ N(0, gain²/fan_in)。")
    print("    左边：第 ℓ 层激活值 A_ℓ 的标准差；右边：从输出端传回一个随机梯度后，第 ℓ 层权重梯度的范数 ‖∂J/∂W_ℓ‖。\n")
    layers = [1, 2, 5, 10, 20]
    head = "".join(f"{'std A' + str(l):>9}" for l in layers) + "".join(f"{'|dW' + str(l) + '|':>10}" for l in [1, 10, 20])
    print(f"{'act':<6}{'gain':>6}" + head)
    Xd = np.random.default_rng(44).normal(size=(500, 100))
    sig = {}
    for act, gain in [("tanh", 0.5), ("tanh", 1.0), ("tanh", 2.0), ("relu", 1.0), ("relu", np.sqrt(2.0)), ("relu", 2.0)]:
        p = init_params([100] + [100] * 20 + [10], np.random.default_rng(45), gain=gain)
        logits, cache = forward(p, Xd, act)
        stds = [float(cache[0][l].std()) for l in layers]
        dlogits = np.random.default_rng(46).normal(size=logits.shape) / len(Xd)
        grads = backward(p, cache, dlogits, act)
        gnorm = [float(np.linalg.norm(grads[2 * (l - 1)])) for l in [1, 10, 20]]
        sig[act, round(gain, 3)] = (stds, gnorm)
        print(f"{act:<6}{gain:>6.3f}" + "".join(f"{v:>9.4f}" for v in stds) + "".join(f"{v:>10.1e}" for v in gnorm))
    t05, t1, t2 = sig["tanh", 0.5], sig["tanh", 1.0], sig["tanh", 2.0]
    r1, rhe, r2 = sig["relu", 1.0], sig["relu", 1.414], sig["relu", 2.0]
    print(f"""
观察 Observations：
  · (a) 全 0 初始化：隐藏层输出全是 tanh(0) = 0，输出层权重的梯度 Aᵀδ = 0；又因为 W₂ = 0，传回隐藏层的误差 δW₂ᵀ 也是 0。
    除了输出偏置，什么都学不动 —— 准确率 {res['all zeros'][1]:.3f}，只是在猜多数类。
  · 全相同的初始化（W = 0.5）：8 个隐藏单元的输入相同、输出相同、收到的梯度也相同，训练 2000 步后仍然一模一样
    （单元之间的最大差 {res['all W = 0.5'][2]:.0e}，隐藏层输出的秩 = {res['all W = 0.5'][3]}）。8 个单元等于 1 个：准确率 {res['all W = 0.5'][1]:.3f}，
    和真正只有 1 个隐藏单元的网络（{acc1[1]:.3f}）一样。随机初始化的唯一目的就是**打破对称性**；它的测试准确率是 {res['random N(0, 1/fan_in)'][1]:.3f}。
  · (b) 尺度太小（tanh，gain 0.5）：每过一层信号缩小约一半，第 10 层的标准差只剩 {t05[0][3]:.4f}，梯度也全都只有 1e−6 量级 —— 网络学不动。
  · gain = 1（tanh）：激活值缓慢衰减（第 20 层 {t1[0][4]:.2f}），但各层的梯度范数基本相同（{t1[1][0]:.2f} / {t1[1][1]:.2f} / {t1[1][2]:.2f}），信号能来回传。
  · 尺度太大（tanh，gain 2）：tanh 饱和（激活值的标准差约 {t2[0][4]:.2f}，大部分在 ±1 附近），越靠近输入的层梯度越大
    （第 1 层 {t2[1][0]:.0f} vs 第 20 层 {t2[1][2]:.1f}）：梯度爆炸，学习率很难选。
  · ReLU 把一半的输入砍成 0，方差每层减半：gain = 1 时第 20 层只剩 {r1[0][4]:.4f}，梯度 {r1[1][0]:.0e}；gain = √2（He 初始化）正好补回来，
    各层的标准差都在 {min(rhe[0]):.2f}~{max(rhe[0]):.2f}；gain = 2 时每层放大 √2 倍，第 20 层到 {r2[0][4]:.0f}。
    规律：Var(Z_ℓ) ≈ gain² · c · Var(Z_{{ℓ−1}})，c 是激活函数“保留”的比例（tanh 在 0 附近约 1，ReLU 是 1/2）；
    让 gain²·c = 1 就是 Xavier / He 初始化的全部想法。BatchNorm / LayerNorm 则是在训练时持续地把尺度拉回来（讲义 §6.3）。""")


# =============================================================================
# 演示 4：backprop 的代价与向量化
# =============================================================================
def demo_cost():
    print_header("演示 4：backprop 为什么高效 —— 反向 ≈ 2 × 前向；数值梯度要 2P 次前向")
    print("只数矩阵乘法里的乘加次数（multiply-accumulate, MAC），batch 大小 n = 64：\n")
    print(f"{'network':<34}{'params P':>11}{'forward MACs':>15}{'backward MACs':>15}{'bwd/fwd':>9}{'fin.diff / backprop':>21}")
    rows = []
    for sizes in ([784, 256, 10], [784, 256, 256, 10], [784, 1024, 1024, 1024, 10], [64] + [512] * 8 + [10]):
        P = sum(a * b + b for a, b in zip(sizes[:-1], sizes[1:]))
        f, b = count_macs(sizes, 64)
        ratio_fd = 2 * P * f / (f + b)                  # 数值梯度：2P 次前向；backprop：1 次前向 + 1 次反向
        rows.append((sizes, P, f, b))
        label = str(sizes) if len(sizes) < 6 else "[64] + [512]*8 + [10]"
        print(f"{label:<34}{P:>11,}{f:>15,}{b:>15,}{b / f:>9.2f}{ratio_fd:>21,.0f}")

    # 向量化：一次矩阵乘法 = 每个样本的外积之和
    rng = np.random.default_rng(47)
    p = init_params([5, 7, 3], rng)
    X = rng.normal(size=(32, 5))
    y = rng.integers(0, 3, size=32)
    _, g_batch, _ = loss_and_grad(p, X, y)
    g_loop = [np.zeros_like(q) for q in p]
    for i in range(32):                                 # 一个样本一个样本地算，再加起来（再除以 n）
        _, gi, _ = loss_and_grad(p, X[i:i + 1], y[i:i + 1])
        g_loop = [a + b / 32 for a, b in zip(g_loop, gi)]
    vec_diff = max(float(np.max(np.abs(a - b))) for a, b in zip(g_batch, g_loop))
    print(f"\n向量化检查：32 个样本逐个反传再平均 vs 一次批量反传，梯度最大差 = {vec_diff:.1e}")
    big = rows[2]
    deep = rows[-1]
    print(f"""
观察 Observations：
  · 每层前向一次矩阵乘法；反向两次 —— ∂J/∂W = Aᵀδ 一次，δW_ℓᵀ（把误差传给上一层）一次；第一层不用往输入传，只要一次。
    所以反向 ≤ 2 × 前向。前三个网络的比例只有 {rows[0][3] / rows[0][2]:.2f}~{big[3] / big[2]:.2f}，因为 784 维输入的第一层占了大部分计算；
    中间层占多数的深网络（最后一行）就是 {deep[3] / deep[2]:.2f}。“一步训练 ≈ 3 次前向”，这就是 CS336 估算训练成本时
    “6 × 参数数 × token 数”的来历（每个 MAC 算 2 FLOPs）。
  · 关键在于：算出**全部** P 个偏导数的代价和 P 无关，只是常数倍的前向。这是 reverse-mode 自动微分的性质：
    输出是一个标量（损失），从输出往回传一遍，就同时得到它对所有输入（参数）的导数。
    forward-mode（或数值差分）每次只能得到一个方向上的导数，需要 P 遍 —— 对 {big[1]:,} 个参数的网络，数值梯度
    比 backprop 贵约 {2 * big[1] * big[2] / (big[2] + big[3]):,.0f} 倍。
  · 代价是**内存**：反向需要前向时每一层的激活值 A_ℓ、Z_ℓ（代码里的 cache）。大模型训练的显存主要花在这里。
  · 向量化：∂J/∂W = Aᵀδ 这一次矩阵乘法，等于把 n 个样本各自的外积 a⁽ⁱ⁾δ⁽ⁱ⁾ᵀ 加起来（差 {vec_diff:.0e}）。
    数学上完全一样，但一次大的矩阵乘法能用上 BLAS / GPU 的并行，比 Python 循环快几个数量级。""")


def main():
    demo_gradient_check()
    demo_xor()
    demo_initialization()
    demo_cost()


if __name__ == "__main__":
    main()