"""
=============================================================================
第 2 讲 · 代码 3：核方法（kernels: feature maps, validity, kernel ridge regression）
=============================================================================

对应讲义：lectures/lecture02_generative_kernels_svm.md 第 4 节

-----------------------------------------------------------------------------
【这个文件要回答的问题】
-----------------------------------------------------------------------------
    演示 1：kernel trick 到底省了什么？
            (a) 多项式核 (xᵀz + c)² 逐位等于显式特征的内积 φ(x)ᵀφ(z)；特征维数 vs 算一次核的代价。
            (b) RBF 核 = 无穷维特征的内积：截断到前 K 维，误差随 K 迅速下降，但永远需要“更多一维”。
            (c) 特征空间里的 LMS（GD）从 θ = 0 出发，θ 永远在 span{φ(x⁽ⁱ⁾)} 里，所以只要 β 和核矩阵 K 就够了。
    演示 2：什么样的函数是合法的核？检查 Gram 矩阵的最小特征值：合法的核 ≥ 0，tanh / 盒子 / 负距离 < 0。
            不合法意味着什么：“特征空间里的长度平方”会是负数。
    演示 3：kernel ridge regression 拟合第 1 讲的正弦数据（同一份 15 个点），RBF 带宽 σ 是复杂度旋钮；
            primal 与 dual 给出同一个预测；O(D³) vs O(n³) —— 什么时候该用核。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 03_kernels.py          # 不到 1 秒

依赖本目录的 kernel_methods.py，以及第 1 讲的 linear_models.py / datasets.py（kernel_methods 导入时自动加入搜索路径，
所以下面要先 import kernel_methods）。
"""

import math

import numpy as np

from kernel_methods import (box_kernel, fit_kernel_ridge, gram_min_eig, kernel_lms, laplacian_kernel,
                            linear_kernel, poly2_features, poly_kernel, rbf_kernel, rbf_taylor_features_1d,
                            sq_dists, tanh_kernel)
from datasets import make_sine
from linear_models import mse, print_header


# =============================================================================
# 演示 1：kernel trick 省了什么
# =============================================================================
def demo_feature_maps():
    print_header("演示 1：核 = 特征空间里的内积，但不用把特征写出来")
    rng = np.random.default_rng(51)

    # (a) 多项式核 vs 显式特征
    X = rng.normal(size=(5, 3))
    Phi = poly2_features(X, c=1.0)
    diff = np.max(np.abs(poly_kernel(X, X, degree=2, c=1.0) - Phi @ Phi.T))
    print(f"(a) d = 3 的 5 个随机点：(xᵀz + 1)² vs φ(x)ᵀφ(z)（φ 有 d² + d + 1 = {Phi.shape[1]} 维），max |diff| = {diff:.1e}")
    print("    特征维数（次数 ≤ p 的不同单项式个数 C(d+p, p)）vs 算一次核的代价（一个 d 维内积）：\n")
    print(f"{'d':>8}{'p':>4}{'feature dim D':>18}{'kernel cost':>14}")
    for d, p in [(3, 2), (10, 2), (100, 2), (100, 3), (100, 5), (1000, 3), (784, 4)]:
        print(f"{d:>8}{p:>4}{math.comb(d + p, p):>18,}{'~' + str(d) + ' mults':>14}")

    # (b) RBF 的显式特征（Taylor 展开）截断
    print("\n(b) 一维 RBF 核（σ = 1）的显式特征 φ_k(x) = e^{−x²/2} x^k / √(k!)，k = 0, 1, 2, …（无穷多维）。")
    print("    只保留前 K 维时，在 [−r, r] 的网格上近似 K(x, z) 的最大误差：\n")
    Ks = [1, 2, 4, 8, 16, 32]
    print(f"{'range':>10}" + "".join(f"{'K=' + str(k):>11}" for k in Ks))
    for r in (1.0, 3.0, 5.0):
        grid = np.linspace(-r, r, 41)
        exact = rbf_kernel(grid, grid, sigma=1.0)
        errs = []
        for k in Ks:
            F = rbf_taylor_features_1d(grid, sigma=1.0, n_terms=k)
            errs.append(np.max(np.abs(exact - F @ F.T)))
        print(f"{'|x| <= ' + format(r, 'g'):>10}" + "".join(f"{e:>11.1e}" for e in errs))

    # (c) 特征空间的 LMS vs 核化的 β 更新
    Xtr = rng.normal(size=(20, 3))
    ytr = Xtr[:, 0] * Xtr[:, 1] - Xtr[:, 2] + 0.1 * rng.normal(size=20)
    Xq = rng.normal(size=(5, 3))
    Phi_tr, Phi_q = poly2_features(Xtr), poly2_features(Xq)
    K = Phi_tr @ Phi_tr.T
    lr = 1.0 / np.linalg.eigvalsh(K)[-1]                  # 保证收敛的步长（第 1 讲：α < 2/λ_max）
    theta = np.zeros(Phi_tr.shape[1])
    for _ in range(500):
        theta = theta + lr * Phi_tr.T @ (ytr - Phi_tr @ theta)    # 显式特征空间里的 batch GD（LMS）
    beta = kernel_lms(poly_kernel(Xtr, Xtr, 2, 1.0), ytr, lr, 500)   # 只用核矩阵
    print("\n(c) 20 个点、13 维特征，GD 500 步：")
    print(f"    显式 θ vs Σᵢ βᵢ φ(x⁽ⁱ⁾)：max |diff| = {np.max(np.abs(theta - Phi_tr.T @ beta)):.1e}")
    print(f"    5 个新点上的预测 θᵀφ(x) vs Σᵢ βᵢ K(x⁽ⁱ⁾, x)：max |diff| = "
          f"{np.max(np.abs(Phi_q @ theta - poly_kernel(Xq, Xtr, 2, 1.0) @ beta)):.1e}")
    print(f"""
观察 Observations：
  · (a) 核函数和显式特征给出同一个内积（差 {diff:.0e}，舍入误差）。但显式特征的维数随 p 爆炸：
    d = 100、p = 5 时约 {math.comb(105, 5) / 1e4:.0f} 万维，而算一次核只要一个 100 维内积加一次求幂。
    （MNIST 的 784 维像素、4 次多项式：约 {math.comb(788, 4) / 1e8:.0f} 亿维 —— 一个样本的特征向量（float64）就要 {math.comb(788, 4) * 8 / 1e9:.0f} GB。）
  · (b) RBF 的特征空间是无穷维的；截断到 K 维的误差随 K 下降得极快（|x| ≤ 1 时 K = 16 已到机器精度），
    但 |x| 越大需要的维数越多（|x| ≤ 5 时 K = 32 还不够）—— 没有任何有限的 K 对所有输入都够用。
    核函数用 O(d) 的代价“算完了”这个无穷级数。
  · (c) 特征空间里的 GD 从 θ = 0 出发，每一步加的都是 φ(x⁽ⁱ⁾) 的线性组合，所以 θ 永远等于 Σᵢ βᵢ φ(x⁽ⁱ⁾)，
    把 θ 换成 β 之后 φ 只通过内积 K(x, z) 出现。这就是 kernel trick 的全部前提：
    算法只用到内积 + 解永远在数据张成的空间里（representer theorem）。""")


# =============================================================================
# 演示 2：合法的核 = Gram 矩阵半正定
# =============================================================================
def demo_validity():
    print_header("演示 2：什么样的“相似度”是合法的核？—— 看 Gram 矩阵的特征值")
    rng = np.random.default_rng(52)
    X = rng.normal(size=(60, 2))
    kernels = [
        ("linear  x'z", linear_kernel(X, X)),
        ("poly    (x'z + 1)^3", poly_kernel(X, X, 3, 1.0)),
        ("RBF     sigma = 1", rbf_kernel(X, X, 1.0)),
        ("Laplace exp(-|x-z|)", laplacian_kernel(X, X, 1.0)),
        ("RBF + poly (sum)", rbf_kernel(X, X, 1.0) + poly_kernel(X, X, 2, 1.0)),
        ("RBF * poly (product)", rbf_kernel(X, X, 1.0) * poly_kernel(X, X, 2, 1.0)),
        ("tanh(x'z - 1)", tanh_kernel(X, X, 1.0, -1.0)),
        ("box 1{|x-z| < 1}", box_kernel(X, X, 1.0)),
        ("-|x-z|^2", -sq_dists(X, X)),
    ]
    print("60 个二维随机点上的 Gram 矩阵 K[i, j] = K(x⁽ⁱ⁾, x⁽ʲ⁾)。合法的核 ⇔ 对任何点集 K 都半正定（Mercer）。")
    print("“负特征值个数”只数 < −1e−10·max|K| 的（更小的是舍入误差）。\n")
    print(f"{'kernel':<24}{'min eig':>12}{'max eig':>12}{'# neg eigs':>12}{'verdict':>10}")
    verdicts = {}
    for name, K in kernels:
        ev = np.linalg.eigvalsh(0.5 * (K + K.T))
        n_neg = int(np.sum(ev < -1e-10 * np.max(np.abs(K))))
        verdicts[name] = (gram_min_eig(K), n_neg)
        print(f"{name:<24}{ev[0]:>12.2e}{ev[-1]:>12.2e}{n_neg:>12}{'valid?' if n_neg == 0 else 'INVALID':>10}")

    Kt = tanh_kernel(X, X, 1.0, -1.0)
    n_negdiag = int(np.sum(np.diag(Kt) < 0))
    Kb = box_kernel(X, X, 1.0)
    ev, V = np.linalg.eigh(Kb)
    v = V[:, 0]
    print(f"\ntanh(x'z − 1)：有 {n_negdiag} / 60 个点的 K(x, x) < 0 —— 如果 K(x, x) = ‖φ(x)‖²，这就是“长度的平方是负数”。")
    print(f"box 核：取最小特征值对应的特征向量 v，vᵀKv = {v @ Kb @ v:.3f} —— 如果 K = ΦΦᵀ，这就是 ‖Σᵢ vᵢ φ(x⁽ⁱ⁾)‖² < 0。")
    print(f"""
观察 Observations：
  · linear、poly、RBF、Laplace 的最小特征值都 ≥ 0（linear 和 poly 的最小特征值是 1e−14～1e−13 级别的“负数”：
    它们的秩只有 2 和 10，剩下的特征值本该是 0，舍入成了正负很小的数）。“valid?” 带问号：
    一组点上没有负特征值只是“没找到反例”；反过来，**一个**负特征值就足以证明不合法。
  · 合法的核做加法、乘法仍然合法（sum、product 两行）—— 这是构造新核的标准工具：K₁ + K₂ 对应把两组特征拼起来，
    K₁K₂ 对应两组特征的所有乘积（张量积）。
  · tanh、box、负距离平方都有明显的负特征值，不是合法的核。box 1{{‖x − z‖ < 1}} 非常符合“相似度”的直觉
    （对称、自己和自己最相似），却不合法：“像相似度”≠“是某个特征空间里的内积”。
  · 不合法的后果：“特征空间”不存在 —— 上面算出了负的“长度平方”。用在 SVM 里，对偶问题不再是凹的
    （没有唯一最优解，SMO 的曲率 η 可能 ≤ 0）；用在 kernel ridge 里，K + λI 可能不可逆或不再对应一个最小化问题。
    实践中：用合法的核（RBF 几乎总是安全的默认选择），或者自己构造时只用加法、乘法、正系数的组合。""")


# =============================================================================
# 演示 3：kernel ridge regression
# =============================================================================
def effective_dof(K, lam):
    """
    Kernel ridge 的有效自由度（effective degrees of freedom）df = tr(K (K + λI)⁻¹) = Σⱼ λⱼ / (λⱼ + λ)。

    【直觉】预测 ŷ = K(K + λI)⁻¹ y = S y，S 叫 smoother matrix（“帽子矩阵”）。tr(S) 衡量模型有多少个“有效参数”：
        λ → 0 且 K 满秩时 df → n（插值每个点）；λ → ∞ 时 df → 0（常数 0）。
        和第 1 讲 LWR 的“有效样本数”是同一类量：都在回答“模型到底有多灵活”。
    """
    ev = np.clip(np.linalg.eigvalsh(0.5 * (K + K.T)), 0.0, None)
    return float(np.sum(ev / (ev + lam)))


def demo_kernel_ridge():
    print_header("演示 3：Kernel ridge regression —— 第 1 讲的正弦数据，RBF 带宽 σ 是复杂度旋钮")
    rng = np.random.default_rng(21)                     # 和第 1 讲代码 02 演示 1 完全同一份数据
    x_tr, y_tr = make_sine(15, rng, noise_std=0.3)
    x_te, y_te = make_sine(2000, rng, noise_std=0.3)
    lam = 0.1
    print("数据：y = sin(2πx) + N(0, 0.3²)，15 个训练点（和第 1 讲多项式实验同一份；那里最好的是 k = 3，test MSE 0.159）。")
    print(f"min_θ Σ(yᵢ − θᵀφ(xᵢ))² + λ‖θ‖²，λ = {lam:g}；解是 α = (K + λI)⁻¹y，预测 f(x) = Σ αᵢ K(xᵢ, x)。\n")
    print(f"{'sigma':>8}{'train MSE':>12}{'test MSE':>12}{'eff. dof':>10}")
    results = []
    for sigma in [0.01, 0.03, 0.05, 0.1, 0.2, 0.3, 0.5, 1.0, 3.0]:
        K = rbf_kernel(x_tr, x_tr, sigma)
        a = fit_kernel_ridge(K, y_tr, lam)
        tr, te = mse(y_tr, K @ a), mse(y_te, rbf_kernel(x_te, x_tr, sigma) @ a)
        results.append((sigma, tr, te))
        print(f"{sigma:>8g}{tr:>12.4f}{te:>12.4f}{effective_dof(K, lam):>10.2f}")
    best = min(results, key=lambda r: r[2])

    print(f"\n固定 σ = {best[0]:g}，改变 λ：")
    print(f"{'lambda':>8}{'train MSE':>12}{'test MSE':>12}{'eff. dof':>10}")
    K = rbf_kernel(x_tr, x_tr, best[0])
    Kte = rbf_kernel(x_te, x_tr, best[0])
    lam_res = {}
    for lm in [1e-6, 1e-3, 1e-2, 0.1, 1.0, 10.0]:
        a = fit_kernel_ridge(K, y_tr, lm)
        lam_res[lm] = (mse(y_tr, K @ a), mse(y_te, Kte @ a), effective_dof(K, lm))
        print(f"{lm:>8g}{lam_res[lm][0]:>12.4f}{lam_res[lm][1]:>12.4f}{lam_res[lm][2]:>10.2f}")

    # primal vs dual：二次多项式核在一维输入上的显式特征只有 3 维
    Phi_tr, Phi_te = poly2_features(x_tr[:, None]), poly2_features(x_te[:, None])
    th = np.linalg.solve(Phi_tr.T @ Phi_tr + lam * np.eye(Phi_tr.shape[1]), Phi_tr.T @ y_tr)   # primal：D × D
    a = fit_kernel_ridge(poly_kernel(x_tr, x_tr, 2, 1.0), y_tr, lam)                          # dual：n × n
    print(f"\nprimal（解 {Phi_tr.shape[1]}×{Phi_tr.shape[1]} 方程组）vs dual（解 15×15 方程组），二次多项式核：")
    print(f"    2000 个测试点上的预测 max |diff| = {np.max(np.abs(Phi_te @ th - poly_kernel(x_te, x_tr, 2, 1.0) @ a)):.1e}")
    print(f"""
观察 Observations：
  · σ 太小（0.01）：每个训练点只影响自己附近一点点，模型在训练点上很准（train MSE ≈ 0），在其他地方几乎预测 0：
    test MSE {results[0][2]:.3f}，有效自由度接近 n = 15 —— 过拟合。σ 太大（3）：所有点“一样相似”，
    曲线被压成接近直线的样子，train / test MSE 都大 —— 欠拟合。σ 和第 1 讲的多项式次数、LWR 的带宽 τ 是同一个旋钮。
  · 最好的 σ = {best[0]:g}：test MSE {best[2]:.3f}，比第 1 讲同一份数据上最好的多项式（0.159）还好一些。
    RBF 核 = 无穷维特征 + ridge 正则化：表达能力无限，靠 λ 和 σ 控制复杂度。
  · λ 是第二个旋钮：λ 越小有效自由度越大（λ = 1e−6 时 {lam_res[1e-6][2]:.1f}；不到 15 是因为 σ = {best[0]:g} 的核矩阵
    有很多比 1e−6 还小的特征值），曲线开始穿过噪声点，test MSE 涨到 {lam_res[1e-6][1]:.2f}；λ 大时曲线被压向 0
    （ridge 把 θ 往 0 拉，在 RBF 特征下就是把 f(x) 往 0 拉），λ = 10 时 test MSE {lam_res[10.0][1]:.3f}。
    第 3 讲会用交叉验证同时选 σ 和 λ。
  · primal 和 dual 的预测一样（舍入误差）：同一个模型，两种解法。primal 解 D × D（D = 特征维数），
    dual 解 n × n。D 小、n 大（普通线性回归）→ primal；D 巨大或无穷（多项式高次、RBF）→ dual。
    代价是 dual 要存 n × n 的核矩阵、解 O(n³)：n = 10⁶ 时核矩阵有 10¹² 个元素（8 TB）。
    大数据上的核方法要靠近似：random features、Nyström、或者干脆换成神经网络（学出特征，第 4 讲）。""")


def main():
    demo_feature_maps()
    demo_validity()
    demo_kernel_ridge()


if __name__ == "__main__":
    main()