"""
=============================================================================
第 3 讲 · 代码 2：Double descent —— 参数比样本还多时，测试误差为什么又降下来了？
=============================================================================

对应讲义：lectures/lecture03_generalization_regularization.md 第 2 节

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
经典的 U 形曲线说：参数越多，variance 越大。可是现代神经网络的参数远多于样本，却泛化得很好。
用"随机 ReLU 特征 + 最小范数最小二乘"这个最简单的过参数化模型，把这件事拆开看：

    演示 1：特征数 p 从 10 扫到 3000（n = 100 固定）。测试误差在 p = n 处出现尖峰，p ≫ n 时又降下来
            —— double descent。同时打印 ‖θ‖ 和特征矩阵的最小奇异值 σ_min：尖峰处 σ_min ≈ 0、‖θ‖ 爆炸。
    演示 2：机制。纯噪声标签 + 高斯特征，此时测试误差 − σ² = ‖θ‖²，而且有封闭公式
            p < n：σ²·p/(n − p − 1)；  p > n：σ²·n/(p − n − 1)。模拟和公式对照；σ_min 和 |√n − √p| 对照。
    演示 3：第二次下降能不能低于第一次下降的谷底？三个不同的真实函数给出三种答案；调好的 ridge 把尖峰整个抹平。
    演示 4：Implicit regularization。p > n 时有无穷多个插值解，GD 选哪一个取决于**从哪里出发**：
            从 0 出发 → 最小范数解（第 1 讲 §2.3）；从随机初始化出发 → 另一个插值解，测试误差更差。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 02_double_descent.py          # 约 16 秒

依赖本目录的 gen_tools.py。
"""

import numpy as np

from gen_tools import (TEACHER_DIM, fit_ridge, make_teacher, print_header, relu_random_features, teacher_target)

N_TRAIN = 100
NOISE_STD = 0.3


def rf_trial(X, y, X_test, f_test, p, rng, lams=()):
    """
    一次随机特征实验：抽一个随机的第一层 W，算特征，拟合最小范数最小二乘（以及若干个 ridge）。

    【参数】
        X, y          ：训练数据 (n, d)、(n,)
        X_test, f_test：测试输入和**无噪声**的真值（测试 MSE 再加上 σ² 就是带噪声标签的期望测试误差）
        p             ：随机特征个数
        rng           ：抽 W 用的随机数生成器
        lams          ：要额外拟合的 ridge λ 列表（可以为空）
    【返回】dict：train、test（最小范数解的训练 / 测试 MSE）、norm（‖θ‖）、smin（Φ 的第 min(n,p) 个奇异值）、
        ridge（列表，每个 λ 的测试 MSE）
    【numpy 小知识】np.linalg.lstsq 在 p > n（欠定）时返回最小范数解；p < n（超定）时返回普通最小二乘解；
        p = n 时方阵可逆就是唯一解。它内部用 SVD，所以三种情形一个函数全包了。
    """
    d = X.shape[1]
    W = rng.normal(0.0, 1.0 / np.sqrt(d), size=(d, p))
    F = relu_random_features(X, W)
    F_test = relu_random_features(X_test, W)
    theta = np.linalg.lstsq(F, y, rcond=None)[0]
    s = np.linalg.svd(F, compute_uv=False)              # 奇异值，从大到小
    out = {
        "train": float(np.mean((F @ theta - y) ** 2)),
        "test": float(np.mean((F_test @ theta - f_test) ** 2)) + NOISE_STD ** 2,
        "norm": float(np.linalg.norm(theta)),
        "smin": float(s[min(F.shape) - 1]),
        "ridge": [float(np.mean((F_test @ fit_ridge(F, y, lam, intercept=False) - f_test) ** 2)) + NOISE_STD ** 2
                  for lam in lams],
    }
    return out


# =============================================================================
# 演示 1：double descent 曲线
# =============================================================================
def demo_curve():
    print_header("演示 1：随机 ReLU 特征 + 最小范数最小二乘，特征数 p 扫过样本数 n = 100")
    rng = np.random.default_rng(41)
    T = 50
    ps = [10, 20, 40, 60, 80, 90, 100, 110, 120, 150, 200, 400, 1000, 3000]
    lam = 1e-2
    X_test = rng.normal(size=(2000, TEACHER_DIM))
    f_test = teacher_target(X_test)
    print(f"数据：x ~ N(0, I_10)，y = uᵀx + ½((vᵀx)² − 1) + N(0, 0.3²)，n = {N_TRAIN}。特征 φ(x) = max(0, Wᵀx)/√p，W 随机且固定。")
    print(f"每个 p 重复 {T} 次（每次新的训练集、新的 W），表中是 {T} 次的平均（test 另给中位数）。")
    print(f"只预测 0 的模型测试误差 = Var f + σ² = 1.5 + 0.09 = 1.59。最后一列是 ridge（λ = {lam:g}）的测试误差。\n")
    res = {p: [] for p in ps}
    for _ in range(T):
        X, y = make_teacher(N_TRAIN, rng, noise_std=NOISE_STD)
        for p in ps:
            res[p].append(rf_trial(X, y, X_test, f_test, p, rng, lams=[lam]))
    print(f"{'p':>6}{'p/n':>6}{'train':>10}{'test(mean)':>12}{'test(median)':>14}{'med |theta|':>13}"
          f"{'med s_min':>11}{'ridge test':>12}")
    summary = {}
    for p in ps:
        tr = np.mean([r["train"] for r in res[p]])
        te = np.array([r["test"] for r in res[p]])
        nrm = np.median([r["norm"] for r in res[p]])
        smin = np.median([r["smin"] for r in res[p]])
        rt = np.mean([r["ridge"][0] for r in res[p]])
        summary[p] = (float(np.mean(te)), float(nrm), float(smin), float(rt))
        print(f"{p:>6}{p / N_TRAIN:>6.1f}{tr:>10.2e}{np.mean(te):>12.4g}{np.median(te):>14.4g}{nrm:>13.4g}"
              f"{smin:>11.3g}{rt:>12.4f}")
    under = [p for p in ps if p < N_TRAIN]
    best_under = min(under, key=lambda p: summary[p][0])
    print(f"""
观察 Observations：
  · 第一次下降（p < n）：经典的 U 形，p = {best_under} 时最好（{summary[best_under][0]:.3f}）；p 接近 n 时测试误差急剧上升。
  · 尖峰在 p = n = {N_TRAIN}：平均测试误差 {summary[N_TRAIN][0]:.3g}。此时恰好有唯一的插值解（方阵），
    最小奇异值 σ_min 的中位数只有 {summary[N_TRAIN][2]:.2g}（p = 1000 时是 {summary[1000][2]:.2g}），‖θ‖ 的中位数 {summary[N_TRAIN][1]:.3g}（p = 1000 时 {summary[1000][1]:.3g}）。
  · 第二次下降（p > n）：训练误差一直是 0（插值），测试误差却随 p 单调下降，p = 3000 时 {summary[3000][0]:.3f}。
    "参数越多越过拟合"在 p > n 之后不成立了 —— 因为 lstsq 在无穷多个插值解里挑了范数最小的那个，
    p 越大，可选的解越多，最小的范数越小（‖θ‖ 从 {summary[N_TRAIN][1]:.3g} 降到 {summary[3000][1]:.3g}）。
  · ridge（λ = 0.01）没有尖峰（p = n 时 {summary[N_TRAIN][3]:.3f}）：尖峰是"恰好插值 + 不加正则"这个组合造成的，
    不是"参数多"本身造成的。但这个 λ 在 p 很大时反而不如不加正则（p = 3000：{summary[3000][3]:.3f} vs {summary[3000][0]:.3f}）——
    λ 太大了，偏差占了上风；λ 要调（演示 3）。演示 2 解释尖峰的机制。""")


# =============================================================================
# 演示 2：机制 —— 纯噪声 + 高斯特征，有封闭公式
# =============================================================================
def demo_mechanism():
    print_header("演示 2：为什么 p ≈ n 最差 —— 最小奇异值和噪声放大")
    rng = np.random.default_rng(42)
    n, sigma, T = 100, 1.0, 400
    ps = [10, 50, 80, 90, 95, 98, 102, 105, 110, 120, 150, 200, 500, 2000]
    print("最干净的设置：标签是纯噪声 y = ε ~ N(0, 1)（真实函数 f = 0），特征 Φ 的元素 iid N(0, 1)。")
    print("测试点 x ~ N(0, I_p)，所以 E[(xᵀθ)²] = ‖θ‖²：测试误差 − σ² 恰好等于 ‖θ‖²，全部是 variance。")
    print("最小范数解 θ = Φ⁺ε = Σᵢ (uᵢᵀε / sᵢ) vᵢ：噪声在第 i 个奇异方向上被放大 1/sᵢ 倍。")
    print(f"每个 p 重复 {T} 次。公式（Gaussian design 的精确期望）：p < n − 1 时 p/(n − p − 1)，p > n + 1 时 n/(p − n − 1)。\n")
    print(f"{'p':>6}{'mean |theta|^2':>16}{'formula':>10}{'median s_min':>14}{'|sqrt(n)-sqrt(p)|':>19}")
    rows = {}
    for p in ps:
        norms, smins = [], []
        for _ in range(T):
            F = rng.normal(size=(n, p))
            eps = rng.normal(0.0, sigma, size=n)
            # 一次 SVD 同时给出 θ 的范数和最小奇异值：‖Φ⁺ε‖² = Σᵢ (uᵢᵀε)² / sᵢ²（只对非零的 sᵢ 求和）。
            # 【numpy 小知识】full_matrices=False 给"瘦"的 SVD：U 是 (n, r)，s 是 (r,)，r = min(n, p)。
            U, sv, _ = np.linalg.svd(F, full_matrices=False)
            norms.append(float(np.sum((U.T @ eps) ** 2 / sv ** 2)))
            smins.append(float(sv[-1]))
        formula = p / (n - p - 1) if p < n - 1 else n / (p - n - 1)
        rows[p] = (float(np.mean(norms)), formula, float(np.median(smins)))
        print(f"{p:>6}{np.mean(norms):>16.4g}{formula * sigma ** 2:>10.4g}{np.median(smins):>14.3g}"
              f"{abs(np.sqrt(n) - np.sqrt(p)):>19.3g}")
    print(f"""
观察 Observations：
  · 模拟和公式对得上（远离 p = n 时误差在几个百分点以内；p = 98、102 时 1/s² 的分布尾巴很重，
    {T} 次的平均也不稳定：{rows[98][0]:.3g} vs {rows[98][1]:.3g}，{rows[102][0]:.3g} vs {rows[102][1]:.3g}）。两个公式在 p → n 时都趋于无穷。
  · 原因在最小奇异值：随机矩阵理论（Marchenko–Pastur / Bai–Yin）说 n × p 高斯矩阵的 σ_min ≈ |√n − √p|，
    p 接近 n 时它贴着 0（p = 98、102 时中位数只有 {rows[98][2]:.2g}，比 |√n − √p| 略大：这个近似是 n、p → ∞ 的渐近结果）。
    噪声在这个方向上被放大 1/σ_min 倍，‖θ‖² 里就多出 (uᵀε)²/σ_min² 这一大项。
  · 直觉：p < n 时，每多一个参数就要用同样多的数据多估一个数（经典的 variance 增长）；
    p = n 时，方程组恰好被"确定"，没有任何余地，必须扭曲 θ 去穿过每一个噪声点；
    p > n 时，插值解有一整个 (p − n) 维的仿射空间，最小范数解可以在里面挑一个"温和"的 ——
    多出来的维度把噪声"摊薄"了（n/(p − n − 1) → 0）。
  · 注意 p > n 时 variance → 0，但如果真实函数 f ≠ 0，bias 会因为"范数最小"的偏好而增大 ——
    第二次下降能降到多低，取决于真实函数和特征的匹配程度（演示 3）。""")


# =============================================================================
# 演示 3：第二次下降 vs 第一次的谷底；调好的 ridge
# =============================================================================
def demo_compare_targets():
    print_header("演示 3：过参数化的极限比经典的最优点好吗？—— 取决于问题")
    rng = np.random.default_rng(43)
    T = 50
    ps = [10, 20, 40, 60, 80, 100, 200, 1000, 3000]
    lams = [1e-4, 1e-3, 1e-2, 1e-1, 1.0]
    X_test = rng.normal(size=(2000, TEACHER_DIM))
    print(f"三个真实函数（gen_tools.teacher_target），n = {N_TRAIN}，σ = 0.3，每个 p 重复 {T} 次。")
    print(f"ridge 的 λ 从 {lams} 里按平均测试误差挑最好的（oracle：直接看测试集，只为说明 ridge 的潜力，实践中要用 CV）。\n")
    print(f"{'target':<10}{'predict 0':>10}{'best p<n':>14}{'p = n':>10}{'p = 3000':>10}{'ridge@p=n':>11}"
          f"{'ridge@3000':>12}{'best ridge':>18}")
    out = {}
    for kind in ["linear", "lin+quad", "tanh"]:
        f_test = teacher_target(X_test, kind)
        res = {p: [] for p in ps}
        for _ in range(T):
            X, y = make_teacher(N_TRAIN, rng, kind=kind, noise_std=NOISE_STD)
            for p in ps:
                res[p].append(rf_trial(X, y, X_test, f_test, p, rng, lams=lams))
        mn = {p: float(np.mean([r["test"] for r in res[p]])) for p in ps}
        rd = {p: np.mean([r["ridge"] for r in res[p]], axis=0) for p in ps}    # (len(lams),)
        best_under = min([p for p in ps if p < N_TRAIN], key=lambda p: mn[p])
        best_r = min(((p, i) for p in ps for i in range(len(lams))), key=lambda pi: rd[pi[0]][pi[1]])
        null = float(np.var(f_test) + NOISE_STD ** 2)
        out[kind] = (null, best_under, mn[best_under], mn[N_TRAIN], mn[3000], float(rd[N_TRAIN].min()),
                     float(rd[3000].min()), best_r, float(rd[best_r[0]][best_r[1]]))
        o = out[kind]
        best_txt = f"{o[8]:.3f} (p={best_r[0]}, λ={lams[best_r[1]]:g})"
        print(f"{kind:<10}{null:>10.3f}{o[2]:>8.3f} (p={best_under:<2}){o[3]:>10.3g}{o[4]:>10.3f}{o[5]:>11.3f}"
              f"{o[6]:>12.3f}{best_txt:>26}")
    lin, lq, th = out["linear"], out["lin+quad"], out["tanh"]
    print(f"""
观察 Observations：
  · 三个函数都有 p = n 的尖峰（{lin[3]:.3g} / {lq[3]:.3g} / {th[3]:.3g}），也都有第二次下降。但第二次下降的终点（p = 3000）和第一次的谷底相比：
      linear   ：{lin[4]:.3f} vs {lin[2]:.3f}
      lin+quad ：{lq[4]:.3f} vs {lq[2]:.3f}
      tanh     ：{th[4]:.3f} vs {th[2]:.3f}
    所以"过参数化的插值解比经典的最优模型更好"不是定理，而是取决于真实函数和特征（以及 n、σ）的匹配程度。
  · 调好 λ 的 ridge 在 p = n 处完全没有尖峰（{lin[5]:.3f} / {lq[5]:.3f} / {th[5]:.3f}），而且三个问题里最好的结果都来自
    "很多特征 + 合适的 ridge"。Nakkiran et al. (2021) 证明了：在某些设置下，最优正则化让测试误差随 p 单调下降。
  · 实践意义：double descent 的尖峰是"不加正则、恰好插值"的产物；现代深度学习几乎总有某种（显式或隐式的）正则化，
    所以很少在实际训练里看到这么夸张的尖峰。但它解释了一个重要事实：参数个数本身不是复杂度的好度量 ——
    ‖θ‖（以及优化器偏好哪种解）才是。""")


# =============================================================================
# 演示 4：implicit regularization —— GD 从哪里出发决定了它找到哪个插值解
# =============================================================================
def gd_least_squares(F, y, theta0, steps):
    """
    在 J(θ) = (1/2n)‖Fθ − y‖² 上做 batch GD，步长 1/λ_max(FᵀF/n)。

    【思路】每一步 θ ← θ − α Fᵀ(Fθ − y)/n。更新量永远在 F 的行空间里（Fᵀ 乘以某个向量），
        所以 θ 在零空间（null space）里的分量**永远不变**：从 θ₀ 出发，GD 的极限是
            θ_∞ = (θ₀ 的零空间分量) + (行空间里唯一的插值解) = (I − F⁺F) θ₀ + F⁺ y
        θ₀ = 0 时就是最小范数解 F⁺y。
    【参数】F：(n, p)；y：(n,)；theta0：(p,) 初始点；steps：步数
    【返回】θ（steps 步之后）
    【边界】行空间里的收敛速度由 FFᵀ/n 的最小非零特征值决定。随机 ReLU 特征全是正数，
        F 有一个很大的"均值方向"奇异值，σ_max/σ_min 可以到 10² 量级，于是 GD 要几十万步才能完全收敛。
        演示里不等它收敛：零空间分量"一步也不动"这件事在每一步都精确成立，极限可以直接用公式算。
    """
    n = len(y)
    lr = 1.0 / np.linalg.eigvalsh(F.T @ F / n)[-1]
    theta = theta0.copy()
    for _ in range(steps):
        theta -= lr * (F.T @ (F @ theta - y)) / n
    return theta


def demo_implicit():
    print_header("演示 4：Implicit regularization —— 同样插值训练集，GD 的起点决定它停在哪个解")
    rng = np.random.default_rng(44)
    p, steps = 400, 20000
    X, y = make_teacher(N_TRAIN, rng, noise_std=NOISE_STD)
    X_test = rng.normal(size=(2000, TEACHER_DIM))
    f_test = teacher_target(X_test)
    W = rng.normal(0.0, 1.0 / np.sqrt(TEACHER_DIM), size=(TEACHER_DIM, p))
    F, F_test = relu_random_features(X, W), relu_random_features(X_test, W)
    F_pinv = np.linalg.pinv(F)
    theta_mn = F_pinv @ y                                        # 最小范数解 = 从 0 出发的 GD 的极限
    sv = np.linalg.svd(F, compute_uv=False)

    def test_mse(theta):
        return float(np.mean((F_test @ theta - f_test) ** 2)) + NOISE_STD ** 2

    direction = rng.normal(size=p) / np.sqrt(p)                  # 一个固定的随机方向（‖·‖ ≈ 1）
    print(f"n = {N_TRAIN}，p = {p}（过参数化 4 倍），同一份数据、同一组特征。初始点 θ₀ = scale × (一个固定的随机单位方向)。")
    print(f"F 的奇异值从 {sv[0]:.2f} 到 {sv[-1]:.4f}：GD（步长 1/λ_max）在最慢的方向上每步只收缩 1 − (σ_min/σ_max)² ≈ 1 − {(sv[-1] / sv[0]) ** 2:.1e}，")
    print(f"所以跑 {steps} 步还没完全收敛。极限 θ_∞ = (I − F⁺F)θ₀ + F⁺y 用公式精确计算。\n")
    print(f"{'init scale':>11}{'|null(theta0)|':>16}{'null drift':>12}{'train@20k':>11}{'test@20k':>10}"
          f"{'test(limit)':>13}{'|theta_limit|':>15}")
    rows = {}
    for scale in [0.0, 5.0, 20.0, 50.0]:
        theta0 = scale * direction
        th = gd_least_squares(F, y, theta0, steps)
        null0 = theta0 - F_pinv @ (F @ theta0)                  # (I − F⁺F) θ₀：θ₀ 的零空间分量
        null_t = th - F_pinv @ (F @ th)                         # 跑完 GD 之后的零空间分量
        drift = float(np.linalg.norm(null_t - null0))           # 理论上恰好是 0
        limit = null0 + theta_mn
        rows[scale] = (float(np.linalg.norm(null0)), drift, test_mse(th), test_mse(limit))
        print(f"{scale:>11g}{rows[scale][0]:>16.3f}{drift:>12.1e}{np.mean((F @ th - y) ** 2):>11.1e}"
              f"{rows[scale][2]:>10.4f}{rows[scale][3]:>13.4f}{np.linalg.norm(limit):>15.3f}")
    print(f"""
观察 Observations：
  · "null drift" 一列是 GD 跑完 {steps} 步之后零空间分量的变化量：在舍入误差的量级（≤ {max(r[1] for r in rows.values()):.0e}）。
    GD 的每一步都是 Fᵀ × (某个向量)，永远在行空间里，零空间分量**一步也不会动**。
  · 所以极限 = 最小范数解 + θ₀ 的零空间分量。四个极限都插值训练集（训练误差 = 0），训练损失完全分不出它们，
    测试误差却随起点变差：{rows[0.0][3]:.4f}（从 0 出发）→ {rows[50.0][3]:.4f}（scale = 50）。
  · 这就是 implicit regularization：我们从没写过"最小化 ‖θ‖"，但"从 0 开始做 GD"这个选择悄悄做了这件事。
    神经网络里也一样 —— 初始化的尺度、优化器（GD / Adam）、步长、batch size 都会影响找到哪个解，
    而训练损失对此一无所知。分类问题里对应的结论是第 1 讲 §3.5：GD 的方向收敛到 max-margin 解。
  · 还有一个附带的现象：从 0 出发跑 {steps} 步的 GD（还没收敛，训练误差不为 0）测试误差 {rows[0.0][2]:.4f}，
    比它的极限 {rows[0.0][3]:.4f} {'更好' if rows[0.0][2] < rows[0.0][3] else '更差'}。GD 先拟合大奇异值的方向、最后才去拟合 σ_min 方向上的噪声 ——
    提前停下来就相当于 ridge。这就是 early stopping（讲义 §4.4，代码 04 的演示 4）。""")


def main():
    demo_curve()
    demo_mechanism()
    demo_compare_targets()
    demo_implicit()


if __name__ == "__main__":
    main()