"""
=============================================================================
第 4 讲 · 代码 5：学出来的特征 vs 随机特征；万能逼近；优化器
=============================================================================

对应讲义：lectures/lecture04_trees_ensembles_nn.md 第 5.3 节、第 6 节

-----------------------------------------------------------------------------
【这个文件要回答的问题】
-----------------------------------------------------------------------------
“神经网络 = 学出来的特征 + GLM 头”。如果特征不学（随机生成后冻结），只训练 GLM 头，会差多少？

    演示 1：万能逼近（universal approximation）的直觉。一维时，h 个 ReLU 单元的网络就是一条有 h 个折点的折线；
            直接**构造**一个插值 sin(2πx) 的网络，误差 ∝ 1/h²。但“存在”不等于“训练得到”：从随机初始化训练同样大小的网络，
            结果并不单调。
    演示 2：学出来的特征 vs 随机特征（random features）。同一个初始化，一个训练全部参数，一个冻结隐藏层、只训练输出头
            （= 在固定特征上做逻辑回归，用第 1 讲的 Newton 法解到最优）。在螺旋数据上比较不同宽度 h 需要多少单元。
            随机特征的效果还强烈依赖于特征的尺度（类似 kernel 的带宽，第 2 讲）。
    演示 3：优化器（简述）。同一个网络，GD / momentum / Adam 各扫 4 个学习率：训练损失下降的速度。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 05_learned_features.py          # 约 11 秒

依赖本目录的 neural_net.py、nonlinear_data.py，以及第 1 讲的 linear_models.py。
"""

import numpy as np

from neural_net import forward, hidden_features, init_params, loss_and_grad, predict, train
from nonlinear_data import make_spirals, sine_target
from linear_models import accuracy, add_intercept, fit_logistic_newton, print_header


# =============================================================================
# 演示 1：万能逼近的直觉 —— ReLU 网络 = 折线
# =============================================================================
def relu_interpolant(f, h):
    """
    构造一个 [1, h, 1] 的 ReLU 网络，它在 [0, 1] 上等于 f 在 h + 1 个等距节点上的分段线性插值。

    【思路】节点 t_k = k/h（k = 0..h）。折线在 [t_k, t_{k+1}] 上的斜率是 s_k = (f(t_{k+1}) − f(t_k))·h。
        写成 ReLU 的和：f̂(x) = f(0) + s_0·ReLU(x − t_0) + Σ_{k≥1} (s_k − s_{k−1})·ReLU(x − t_k)。
        每个 ReLU(x − t_k) 在 t_k 处"打开"一个新的斜率变化 —— 第 k 个隐藏单元就是第 k 个折点。
        对应的网络参数：W₁ = 1（全部），b₁ = −t_k，W₂ = 斜率的变化量，b₂ = f(0)。
    【参数】f：一元函数（接受 numpy 数组）；h：隐藏单元数（≥ 1）
    【返回】params 列表 [W1, b1, W2, b2]，可以直接交给 neural_net.forward(…, act="relu")
    【ML 联系】这就是万能逼近定理在一维的构造性证明：任何连续函数都能被折线一致逼近，而折线 = ReLU 网络。
        对光滑的 f，插值误差 ≤ max|f''|/(8h²)，所以误差 ∝ 1/h²。高维时同样的构造需要的单元数随维数指数增长 ——
        定理只保证“存在”，不保证“单元数合理”，更不保证“GD 找得到”。
    """
    t = np.linspace(0.0, 1.0, h + 1)
    s = np.diff(f(t)) * h                                  # h 段的斜率
    dslope = np.concatenate([[s[0]], np.diff(s)])          # 每个折点处的斜率变化
    W1 = np.ones((1, h))
    b1 = -t[:-1]
    W2 = dslope[:, None]
    b2 = np.array([f(np.array([0.0]))[0]])
    return [W1, b1, W2, b2]


def demo_universal_approximation():
    print_header("演示 1：万能逼近的直觉 —— 一维 ReLU 网络就是折线")
    x = np.linspace(0.0, 1.0, 2001)
    fx = sine_target(x)
    print("目标 f(x) = sin(2πx)，x ∈ [0, 1]（无噪声，只看“能不能表示”）。")
    print("constructed：直接写出插值 h + 1 个等距节点的 ReLU 网络（不训练）；")
    print("trained：同样大小的网络从随机初始化开始，用 Adam（lr = 0.01）在 200 个等距点上训练 4000 步。\n")
    print(f"{'h':>4}{'constructed max err':>21}{'max err * h^2':>15}{'trained max err':>17}{'trained: kinks in [0,1]':>25}")
    xs = np.linspace(0.0, 1.0, 200)
    res = {}
    for h in [2, 4, 8, 16, 32, 64]:
        p = relu_interpolant(sine_target, h)
        err_c = float(np.max(np.abs(forward(p, x[:, None], "relu")[0][:, 0] - fx)))
        rng = np.random.default_rng(51)
        q = init_params([1, h, 1], rng)
        q[0] = rng.normal(0.0, 3.0, size=(1, h))           # 折点 −b/w 大致散布在 [−1, 2] 上
        q[1] = rng.uniform(-3.0, 3.0, size=h)
        q = train(q, xs[:, None], sine_target(xs), steps=4000, lr=0.01, act="relu", head="mse")
        err_t = float(np.max(np.abs(forward(q, x[:, None], "relu")[0][:, 0] - fx)))
        with np.errstate(divide="ignore", invalid="ignore"):  # w = 0 的单元没有折点（−b/0 = ±inf），直接不计
            kinks = -q[1] / q[0][0]
        n_k = int(np.sum((kinks > 0) & (kinks < 1)))
        res[h] = (err_c, err_t, n_k)
        print(f"{h:>4}{err_c:>21.5f}{err_c * h * h:>15.3f}{err_t:>17.4f}{n_k:>25}")
    worse = [h for h in res if h > 2 and res[h][1] > res[h // 2][1]]
    if worse:
        monotone_note = "而且不单调（" + "、".join(f"h = {h}" for h in worse) + " 比一半宽度的网络还差）。单元也用得很浪费："
    else:
        monotone_note = "这次误差随 h 单调下降，但单元用得很浪费："
    print(f"""
观察 Observations：
  · 构造出来的网络误差按 1/h² 下降：误差 × h² 随 h 增大趋于常数（h = 64 时 {res[64][0] * 64 * 64:.2f}，理论上界 max|f''|/8 = (2π)²/8 ≈ 4.93）。
    这就是万能逼近定理（Cybenko 1989；Hornik 1991）在一维的样子：足够多的单元，就能把任何连续函数逼近到任意精度。
  · 但定理只说“存在”。训练出来的网络：h = 64 时最大误差 {res[64][1]:.3f}，比同样大小的构造解（{res[64][0]:.5f}）差一个数量级以上；
    {monotone_note}很多单元的折点落在 [0, 1] 之外（对拟合毫无贡献），
    落在 [0, 1] 里的折点数（最后一列）远少于 h，h = 4 时只有 {res[4][2]} 个。这是单个随机种子的结果，有运气成分；
    但和代码 4 的 XOR 一样，传达的是同一件事：**能表示 ≠ 训练得到**。
  · In high dimensions, "placing kinks on a grid" needs a number of units exponential in the dimension (the curse of
    dimensionality). Networks work in practice not because of this theorem, but because real data have structure and
    training moves the units to where they are needed (demo 2).""")


# =============================================================================
# 演示 2：学出来的特征 vs 随机特征
# =============================================================================
def spiral_init(h, seed, scale):
    """
    [2, h, 2] 网络的初始化：W₁、b₁ 的每个元素 ~ N(0, scale²)；输出层用 init_params 的默认值。

    【为什么不用 init_params 的 1/fan_in 尺度？】输入只有 2 维，按 1/fan_in 初始化的特征 tanh(w·x + b)
        在 [−1, 1]² 上变化很平缓；螺旋相邻两圈只隔约 0.33，需要更“陡”的特征。scale 就是特征的陡峭程度，
        作用和 RBF kernel 的带宽（第 2 讲）类似。b₁ 也随机，否则所有特征的“等高线”都经过原点。
    """
    rng = np.random.default_rng(seed)
    p = init_params([2, h, 2], rng)
    p[0] = rng.normal(0.0, scale, size=(2, h))
    p[1] = rng.normal(0.0, scale, size=h)
    return p


def random_features_fit(p, X, y, Xt, yt, l2=1e-4):
    """
    随机特征模型：冻结隐藏层 (W₁, b₁)，只在特征 φ(x) = tanh(W₁ᵀx + b₁) 上拟合输出头。

    【思路】K = 2 时 softmax 头和逻辑回归等价，而逻辑回归是凸的：直接用第 1 讲的 fit_logistic_newton 解到最优，
        给随机特征“最好的机会”（不存在训练不充分的问题）。h ≥ n 时数据在特征空间里线性可分，MLE 不存在，
        所以加一点 L2（λ = 1e−4）。
    【返回】(train_acc, test_acc)
    """
    H, Ht = np.tanh(X @ p[0] + p[1]), np.tanh(Xt @ p[0] + p[1])
    theta, _ = fit_logistic_newton(add_intercept(H), y, l2=l2)
    return (accuracy(y, (add_intercept(H) @ theta >= 0).astype(int)),
            accuracy(yt, (add_intercept(Ht) @ theta >= 0).astype(int)))


def demo_learned_vs_random():
    print_header("演示 2：学出来的特征 vs 随机特征（random features）—— 螺旋数据")
    rng = np.random.default_rng(52)
    X, y = make_spirals(200, rng)
    Xt, yt = make_spirals(1000, rng)
    print("两条螺旋线（各转 1.5 圈），训练 400 个点，测试 2000 个点。网络 [2, h, 2]，tanh。")
    print("MLP：训练全部参数（Adam，lr = 0.02，2000 步）。RF：同一个初始化，冻结隐藏层，输出头用 Newton 解到最优。")
    print("“hidden moved”= 训练后隐藏层参数 (W₁, b₁) 的相对变化 ‖Δ‖/‖初始值‖。\n")
    print(f"{'h':>5}{'MLP train':>11}{'MLP test':>10}{'hidden moved':>14}{'RF(s=3) train':>15}{'RF(s=3) test':>14}"
          f"{'RF(s=1) test':>14}")
    res = {}
    for h in [5, 10, 20, 50, 200, 1000]:
        p0 = spiral_init(h, 53, 3.0)
        rf3 = random_features_fit(p0, X, y, Xt, yt)
        rf1 = random_features_fit(spiral_init(h, 53, 1.0), X, y, Xt, yt)
        p = train([a.copy() for a in p0], X, y, steps=2000, lr=0.02)
        moved = float(np.linalg.norm(np.concatenate([(p[0] - p0[0]).ravel(), p[1] - p0[1]]))
                      / np.linalg.norm(np.concatenate([p0[0].ravel(), p0[1]])))
        res[h] = (accuracy(y, predict(p, X)), accuracy(yt, predict(p, Xt)), moved, rf3[0], rf3[1], rf1[1])
        r = res[h]
        print(f"{h:>5}{r[0]:>11.3f}{r[1]:>10.3f}{r[2]:>14.3f}{r[3]:>15.3f}{r[4]:>14.3f}{r[5]:>14.3f}")
    first_mlp = min(h for h in res if res[h][1] >= 0.9)
    first_rf3 = min((h for h in res if res[h][4] >= 0.9), default=None)
    first_rf1 = min((h for h in res if res[h][5] >= 0.9), default=None)
    print(f"""
观察 Observations：
  · 测试准确率第一次达到 0.9 需要的宽度：MLP h = {first_mlp}；随机特征（s = 3）h = {first_rf3}；随机特征（s = 1）h = {first_rf1}。
    学习特征的网络用少得多的单元就够了：它把每个单元的“等高线”挪到螺旋需要的地方，而随机特征只能靠数量去碰运气。
  · 随机特征不是没用：h 很大时（s = 3，h = 1000：{res[1000][4]:.3f}）它和 MLP（{res[1000][1]:.3f}）不相上下 —— 宽度 → ∞ 时，
    随机特征上的线性模型收敛到一个 kernel 方法（Rahimi & Recht 2007；第 2 讲的 kernel trick 反过来用）。
    但它强烈依赖于随机分布的尺度：s = 1 的特征太平缓，同样 h = 200 只有 {res[200][5]:.3f}（s = 3 是 {res[200][4]:.3f}）。
    选 s 就像选 RBF kernel 的带宽；MLP 则会自己把特征调到合适的陡峭程度。
  · “hidden moved”总体随宽度减小（h ≤ 50 时约 {min(res[h][2] for h in (5, 10, 20, 50)):.1f}~{max(res[h][2] for h in (5, 10, 20, 50)):.1f}，h = 200 时 {res[200][2]:.2f}，h = 1000 时 {res[1000][2]:.2f}）：网络越宽，每个隐藏单元需要移动得越少，
    越接近“随机特征 + 训练输出头”。这就是 neural tangent kernel（NTK）/ lazy training 的观察：无限宽的网络在某些
    参数化下等价于一个 kernel 方法。而真正“学特征”（feature learning）的好处，恰恰出现在有限宽度、参数移动较大时。
  · 第 3 讲用的正是随机特征模型来演示 double descent：h ≈ n 时最小范数解的测试误差出现尖峰。""")


# =============================================================================
# 演示 3：优化器
# =============================================================================
def demo_optimizers():
    print_header("演示 3：优化器（简述）—— GD vs momentum vs Adam")
    rng = np.random.default_rng(54)
    X, y = make_spirals(200, rng)
    Xt, yt = make_spirals(1000, rng)
    print("同一个 [2, 50, 2] tanh 网络（同一个初始化，s = 3），full-batch，3000 步。每个优化器扫 4 个学习率；")
    print("表中是训练集上的 cross-entropy。“best”标出每个优化器在第 3000 步训练损失最低的学习率。\n")
    print(f"{'optimizer':<10}{'lr':>7}{'loss@10':>9}{'loss@100':>10}{'loss@300':>10}{'loss@1000':>11}{'loss@3000':>11}"
          f"{'test acc':>10}")
    grid = {"gd": [0.1, 0.3, 1.0, 3.0], "momentum": [0.01, 0.03, 0.1, 0.3], "adam": [0.003, 0.01, 0.03, 0.1]}
    best, all_runs = {}, {}
    for opt, lrs in grid.items():
        runs = []
        for lr in lrs:
            hist = []
            p = train(spiral_init(50, 55, 3.0), X, y, steps=3000, lr=lr, optimizer=opt, history=hist, log_every=10)
            d = dict(hist)
            runs.append((lr, d, accuracy(yt, predict(p, Xt))))
        b = min(runs, key=lambda r: r[1][3000])
        best[opt], all_runs[opt] = b, runs
        for lr, d, acc in runs:
            mark = "  <- best" if lr == b[0] else ""
            print(f"{opt:<10}{lr:>7g}{d[10]:>9.3f}{d[100]:>10.3f}{d[300]:>10.3f}{d[1000]:>11.4f}{d[3000]:>11.4f}{acc:>10.3f}{mark}")
    g, m, a = best["gd"], best["momentum"], best["adam"]
    # 公平一点的比较：每个检查点上，各优化器在 4 个学习率里最低的训练损失
    env = {opt: {c: min(d[c] for _, d, _ in all_runs[opt]) for c in (100, 300, 1000, 3000)} for opt in grid}
    print("\n每个检查点上、各优化器在 4 个学习率中最低的训练损失（lower envelope）：\n")
    print(f"{'optimizer':<10}{'@100':>9}{'@300':>9}{'@1000':>9}{'@3000':>9}")
    for opt in grid:
        print(f"{opt:<10}" + "".join(f"{env[opt][c]:>9.4f}" for c in (100, 300, 1000, 3000)))
    print(f"""
观察 Observations：
  · 看下包络（每个检查点挑最好的学习率）：第 300 步 GD {env['gd'][300]:.3f}、momentum {env['momentum'][300]:.3f}、Adam {env['adam'][300]:.3f}；
    第 3000 步 {env['gd'][3000]:.4f}、{env['momentum'][3000]:.4f}、{env['adam'][3000]:.4f}。Adam 在每个检查点都最低；momentum 前期明显快于 GD，
    到第 3000 步 GD（用最大的 lr = {g[0]:g}）基本追平了 momentum。
  · GD 的学习率很难选：lr = {g[0]:g} 最后最好，但第 10 步的损失是 {g[1][10]:.2f}（第 1 步是 {g[1][1]:.2f}）—— 一开始步长超过了最陡方向的
    稳定上限 2/λ_max（第 1 讲 §2.2），损失先冲高，等网络走到曲率更小的区域才回落（“catapult / edge of stability”现象）。
    神经网络的损失面同样有病态的条件数，只是曲率还随位置变化。
  · momentum：把过去的梯度累加成“速度”，在一致的方向上加速、在来回震荡的方向上抵消；有效步长约是 lr/(1 − β) = 10·lr，
    所以它的最佳 lr 比 GD 小一个数量级。Adam：每个坐标除以自己梯度大小的滑动平均（√v），相当于逐坐标的 feature scaling，
    对条件数不敏感，所以成了深度学习的默认选择。
  · Lower training loss is not better test accuracy: GD {g[2]:.3f}, momentum {m[2]:.3f}, Adam {a[2]:.3f}. The optimizer
    decides how fast we reach low loss; generalization also depends on where we stop (implicit regularization,
    Lecture 3). For warmup / decay schedules, AdamW and normalization layers see CS224N / CS336.""")


def main():
    demo_universal_approximation()
    demo_learned_vs_random()
    demo_optimizers()


if __name__ == "__main__":
    main()
