"""
=============================================================================
第 1 讲 · 代码 2：特征映射与局部加权回归（feature maps & locally weighted regression）
=============================================================================

对应讲义：lectures/lecture01_linear_models.md 第 2.5 节

-----------------------------------------------------------------------------
【这个文件要回答的问题】
-----------------------------------------------------------------------------
"线性回归"只能拟合直线吗？不是 —— 它只要求模型对 **参数 θ** 是线性的。两种让它变灵活的办法：

    演示 1：Parametric 的办法 —— 换特征。x → φ(x) = [1, x, x², …]。次数太低欠拟合（underfit），
            太高过拟合（overfit）：训练误差一路下降，测试误差先降后升。（第 3 讲 bias–variance 的预告）
    演示 2：Non-parametric 的办法 —— LWR。不选特征，而是对每个查询点只用"附近"的数据拟合一条直线。
            带宽 τ 扮演和"多项式次数"相同的角色：τ 小 → 过拟合，τ 大 → 退化成一条全局直线。
    演示 3：把曲线打印出来看：τ 大小不同的 LWR 长什么样；以及为什么用"局部直线"而不是"局部平均"
            （Nadaraya–Watson）—— 边界处的偏差（boundary bias）。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 02_features_lwr.py          # 约 1 秒

依赖本目录的 datasets.py 和 linear_models.py。
"""

import numpy as np

from datasets import make_sine, sine_target
from linear_models import add_intercept, fit_least_squares, lwr_predict, mse, poly_features, print_header


def to_unit_interval(x):
    """
    把 x ∈ [0, 1] 线性映射到 [−1, 1]。

    【为什么】x ∈ [0, 1] 时 x¹⁴ 的值非常小（0.5¹⁴ ≈ 6e−5），各列的数量级相差巨大，XᵀX 严重病态。
        映射到以 0 为中心的 [−1, 1] 能缓解很多（更彻底的办法是用正交多项式，比如 Legendre / Chebyshev）。
        这只是换坐标，不改变模型能表示的函数集合。
    """
    return 2.0 * np.asarray(x) - 1.0


# =============================================================================
# 演示 1：多项式特征 —— underfit vs overfit
# =============================================================================
def demo_polynomial():
    print_header("演示 1：多项式特征 φ(x) = [1, x, …, x^k] —— 次数 k 从 0 到 14")
    rng = np.random.default_rng(21)
    x_tr, y_tr = make_sine(15, rng, noise_std=0.3)
    x_te, y_te = make_sine(2000, rng, noise_std=0.3)
    print("数据：y = sin(2πx) + N(0, 0.3²)，训练集只有 15 个点，测试集 2000 个点。")
    print("噪声方差 0.09 是测试 MSE 的下限（irreducible error）：再好的模型也预测不了噪声。\n")
    print(f"{'degree k':>9}{'#params':>9}{'train MSE':>12}{'test MSE':>12}{'‖θ‖':>12}")
    for k in [0, 1, 3, 5, 7, 9, 12, 14]:
        Phi_tr = poly_features(to_unit_interval(x_tr), k)
        Phi_te = poly_features(to_unit_interval(x_te), k)
        th = fit_least_squares(Phi_tr, y_tr)
        print(f"{k:>9}{k + 1:>9}{mse(y_tr, Phi_tr @ th):>12.4f}{mse(y_te, Phi_te @ th):>12.4g}"
              f"{np.linalg.norm(th):>12.4g}")
    print("""
观察 Observations：
  · k = 0、1：模型太简单，连训练集都拟合不好（train MSE 远大于 0.09）—— 这是 underfitting / high bias。
  · k = 3、5 最好：测试 MSE ≈ 0.16~0.18。离噪声下限 0.09 还有距离，因为只有 15 个点，参数本身估不准（方差）。
    k = 3 的训练 MSE（0.036）反而低于 0.09：它已经开始拟合一点噪声了 —— 训练误差总是偏乐观。
  · k 继续增大：训练误差一路下降，k = 14 时 15 个参数穿过 15 个点，训练误差 = 0；
    测试误差却爆炸 —— 这是 overfitting / high variance：模型把噪声当成了规律。
  · 注意 ‖θ‖ 的变化：过拟合时系数变得巨大、正负交替，曲线在数据点之间剧烈摆动。
    "惩罚 ‖θ‖ 大的解"（ridge / weight decay）是第 3 讲正则化的直觉来源。
  · 训练误差永远随模型复杂度单调下降（嵌套的模型族），所以**不能用训练误差选模型**。""")


# =============================================================================
# 演示 2：LWR 的带宽 τ
# =============================================================================
def demo_lwr_bandwidth():
    print_header("演示 2：Locally weighted regression —— 带宽 τ 是复杂度旋钮")
    rng = np.random.default_rng(22)
    x_tr, y_tr = make_sine(50, rng, noise_std=0.3)
    x_te, y_te = make_sine(2000, rng, noise_std=0.3)
    print("训练集 50 个点。对每个测试点 x_q，用权重 w = exp(−(x − x_q)²/(2τ²)) 拟合一条直线，取它在 x_q 处的值。\n")
    print(f"{'tau':>8}{'train MSE':>12}{'test MSE':>12}{'有效样本数 (平均)':>18}")
    for tau in [0.003, 0.01, 0.03, 0.06, 0.1, 0.2, 0.5, 10.0]:
        p_tr = lwr_predict(x_tr, y_tr, x_tr, tau)
        p_te = lwr_predict(x_tr, y_tr, x_te, tau)
        # 有效样本数（effective sample size）：(Σw)² / Σw²。权重全相等时等于 n，只有一个非零时等于 1。
        ess = []
        for xq in x_te[:200]:
            w = np.exp(-(x_tr - xq) ** 2 / (2 * tau ** 2))
            ess.append(w.sum() ** 2 / max((w ** 2).sum(), 1e-300))
        print(f"{tau:>8g}{mse(y_tr, p_tr):>12.4f}{mse(y_te, p_te):>12.4f}{np.mean(ess):>18.1f}")
    th = fit_least_squares(add_intercept(x_tr), y_tr)
    print(f"{'global line':>8}{mse(y_tr, add_intercept(x_tr) @ th):>12.4f}{mse(y_te, add_intercept(x_te) @ th):>12.4f}"
          f"{50:>18}")
    print("""
观察 Observations：
  · τ 很小：每个查询点只看最近的一两个点 —— 训练误差接近 0（每个训练点几乎只看自己），测试误差大：overfit。
  · τ 很大：所有点权重几乎一样，LWR 退化成全局最小二乘直线（最后两行几乎相同）：underfit。
  · 中间的 τ 最好。τ 和演示 1 的多项式次数是同一个旋钮的两种形式：模型复杂度 / bias–variance 权衡。
  · "有效样本数"量化了 τ 的含义：每次预测实际上用了多少个点。用的点越少，预测越"看运气"（方差越大）。
  · Parametric vs non-parametric：多项式拟合完只需存 k+1 个数；LWR 必须一直保留全部训练数据，
    每次预测都要重新解一个加权回归（O(n) 起步）。数据多、维度低时 LWR 很好用；维度高时"附近"的点很少
    （curse of dimensionality），它就不灵了。""")


# =============================================================================
# 演示 3：看曲线；局部直线 vs 局部平均
# =============================================================================
def nadaraya_watson(x_train, y_train, x_query, tau):
    """
    Nadaraya–Watson kernel regression（"局部平均"，local constant）：ŷ(x_q) = Σ wᵢ yᵢ / Σ wᵢ。

    【和 LWR 的关系】LWR 在每个查询点拟合一条直线（local linear）；NW 只拟合一个常数（加权平均）。
        在数据的边界上，查询点一侧没有数据，加权平均会被"另一侧"的点拉偏 —— 这叫 boundary bias。
        局部直线能沿着斜率外推一点点，偏差小得多。
    【ML 联系】把 wᵢ 换成 softmax(q·kᵢ)，这就是 attention：ŷ = Σ softmax(q·k)ᵢ vᵢ。
    """
    preds = np.empty(len(x_query))
    for q, xq in enumerate(x_query):
        sq = (x_train - xq) ** 2
        w = np.exp(-(sq - sq.min()) / (2 * tau ** 2))
        preds[q] = np.sum(w * y_train) / np.sum(w)
    return preds


def demo_curves():
    print_header("演示 3：把 LWR 的曲线打印出来；局部直线 vs 局部平均（boundary bias）")
    rng = np.random.default_rng(22)                         # 和演示 2 同一份训练数据
    x_tr, y_tr = make_sine(50, rng, noise_std=0.3)
    grid = np.linspace(0.0, 1.0, 11)
    cols = {
        "true f(x)": sine_target(grid),
        "LWR τ=0.01": lwr_predict(x_tr, y_tr, grid, 0.01),
        "LWR τ=0.06": lwr_predict(x_tr, y_tr, grid, 0.06),
        "LWR τ=10": lwr_predict(x_tr, y_tr, grid, 10.0),
        "NW τ=0.06": nadaraya_watson(x_tr, y_tr, grid, 0.06),
    }
    print(f"{'x':>6}" + "".join(f"{name:>13}" for name in cols))
    for r, xv in enumerate(grid):
        print(f"{xv:>6.1f}" + "".join(f"{v[r]:>13.3f}" for v in cols.values()))
    err = {name: np.abs(v - cols["true f(x)"]) for name, v in cols.items() if name != "true f(x)"}
    print(f"\n  端点 x = 0 和 x = 1 处的平均 |误差|：LWR τ=0.06 = {np.mean(err['LWR τ=0.06'][[0, -1]]):.3f}，"
          f"NW τ=0.06 = {np.mean(err['NW τ=0.06'][[0, -1]]):.3f}")
    print(f"  内部 9 个点的平均 |误差|：        LWR τ=0.06 = {np.mean(err['LWR τ=0.06'][1:-1]):.3f}，"
          f"NW τ=0.06 = {np.mean(err['NW τ=0.06'][1:-1]):.3f}")
    print("""
观察 Observations：
  · τ = 0.01 的曲线跟着单个噪声点上下跳；τ = 10 是一条直线（斜率为负的那条最小二乘直线）；τ = 0.06 最接近真实函数。
  · sin(2πx) 在 x = 0 处向上走、在 x = 1 处从下往上回到 0。在端点，查询点只有一侧有数据：
    NW（局部平均）只能取"内侧"那些点的平均，被系统性地拉向内侧的值；LWR（局部直线）顺着斜率外推，偏差小很多。
    这就是为什么 CS229 讲的是"局部加权**线性**回归"而不是"局部加权平均"。
  · 内部的点上两者一样好（都是 0.070）：那里两侧都有数据，加权平均的偏差能左右抵消。
  · 局部直线也有代价：τ = 0.01 在 x = 1 处预测出 11.4！端点附近只有一两个有效的点，穿过两个靠得很近的点的直线
    斜率可能非常大，外推一点点就飞出去了。τ 太小时，"局部直线"比"局部平均"更不稳定（方差更大）。""")


def main():
    demo_polynomial()
    demo_lwr_bandwidth()
    demo_curves()


if __name__ == "__main__":
    main()