import numpy as np
import matplotlib.pyplot as plt
import matplotlib
from manim import *

# ── 中文字体设置 ──
# 优先使用霞鹜文楷（更美观），回退到 Noto Sans CJK SC
_cjk_fonts = ['LXGW WenKai', 'Noto Sans CJK SC', 'HYWenHei']
_available = [f.name for f in matplotlib.font_manager.fontManager.ttflist]
_cjk = next((f for f in _cjk_fonts if f in _available), None)
if _cjk:
    plt.rcParams['font.sans-serif'] = [_cjk, 'DejaVu Sans']
    plt.rcParams['axes.unicode_minus'] = False
    print(f"中文字体: {_cjk}")
else:
    print("⚠️ 未找到中文字体，图表中的中文可能无法正常显示")

# 全局设置
plt.rcParams.update({
    'figure.figsize': (10, 6),
    'font.size': 12,
    'axes.labelsize': 14,
    'axes.titlesize': 16,
})
np.random.seed(42)

print("NumPy version:", np.__version__)

# 真实参数
TRUE_W = 3.0
TRUE_B = 5.0
NOISE_STD = 4.0

# 生成数据
N = 100
X = np.random.uniform(-5, 5, N)
noise = np.random.normal(0, NOISE_STD, N)
y = TRUE_W * X + TRUE_B + noise

# 可视化
plt.figure(figsize=(8, 5))
plt.scatter(X, y, alpha=0.6, label='带噪声的数据点')
x_line = np.linspace(-6, 6, 100)
plt.plot(x_line, TRUE_W * x_line + TRUE_B, 'r--', linewidth=2, label=f'真实直线: y = {TRUE_W}x + {TRUE_B}')
plt.xlabel('x')
plt.ylabel('y')
plt.legend()
plt.title('带噪声的线性数据')
plt.grid(True, alpha=0.3)
plt.show()

print(f"数据量: N = {N}")
print(f"真实模型: y = {TRUE_W}x + {TRUE_B} + ε,  ε ~ N(0, {NOISE_STD}²)")

def compute_loss(X, y, w, b):
    """计算 MSE 损失"""
    y_pred = w * X + b
    return np.mean((y_pred - y) ** 2)


def compute_gradient(X, y, w, b):
    """计算 w 和 b 的梯度"""
    N = len(X)
    y_pred = w * X + b
    error = y_pred - y
    dw = (2 / N) * np.sum(error * X)
    db = (2 / N) * np.sum(error)
    return dw, db


def gradient_descent(X, y, w_init, b_init, lr, epochs, verbose=True):
    """梯度下降优化，返回权重历史与损失历史"""
    w, b = w_init, b_init
    w_history, b_history, loss_history = [w], [b], [compute_loss(X, y, w, b)]

    for epoch in range(epochs):
        dw, db = compute_gradient(X, y, w, b)
        w -= lr * dw
        b -= lr * db

        w_history.append(w)
        b_history.append(b)
        loss_history.append(compute_loss(X, y, w, b))

        if verbose and epoch % max(1, epochs // 10) == 0:
            print(f"Epoch {epoch:4d}: loss={loss_history[-1]:.4f}, w={w:.4f}, b={b:.4f}")

    return np.array(w_history), np.array(b_history), np.array(loss_history)

# 初始化与训练
w_init, b_init = 0.0, 0.0
lr = 0.01
epochs = 500

w_hist, b_hist, loss_hist = gradient_descent(X, y, w_init, b_init, lr, epochs)

print(f"\n训练完毕！")
print(f"学习到的参数: w = {w_hist[-1]:.4f}（真实值: {TRUE_W}）, b = {b_hist[-1]:.4f}（真实值: {TRUE_B}）")
print(f"最终损失: {loss_hist[-1]:.4f}")

# 可视化：最终拟合结果
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# 左图：数据 + 拟合直线
ax = axes[0]
ax.scatter(X, y, alpha=0.5, label='数据点')
x_line = np.linspace(X.min() - 1, X.max() + 1, 100)
ax.plot(x_line, TRUE_W * x_line + TRUE_B, 'r--', linewidth=2, label=f'真实: y={TRUE_W}x+{TRUE_B}')
ax.plot(x_line, w_hist[-1] * x_line + b_hist[-1], 'g-', linewidth=2, label=f'拟合: y={w_hist[-1]:.2f}x+{b_hist[-1]:.2f}')
ax.set_xlabel('x'); ax.set_ylabel('y')
ax.legend(); ax.set_title('数据与拟合直线')
ax.grid(True, alpha=0.3)

# 右图：损失曲线
ax = axes[1]
ax.plot(loss_hist, 'b-', linewidth=1.5)
ax.set_xlabel('Epoch'); ax.set_ylabel('MSE Loss')
ax.set_title('训练损失曲线')
ax.set_yscale('log')
ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

%%manim -qm -r 960,540 -v WARNING GradientDescentTrajectory

config.media_embed = True
config.media_width = "100%"
config.verbosity = "WARNING"

class GradientDescentTrajectory(Scene):
    def construct(self):
        step = max(1, len(w_hist) // 100)
        ws = w_hist[::step]
        bs = b_hist[::step]
        losses = loss_hist[::step]

        x_min, x_max = X.min() - 1, X.max() + 1
        y_min, y_max = y.min() - 3, y.max() + 3
        axes = Axes(x_range=[x_min, x_max, 2], y_range=[y_min, y_max, 5],
                     x_length=8, y_length=5.5, tips=False).shift(DOWN * 0.3)
        self.add(axes)

        dots = VGroup(*[Dot(axes.coords_to_point(xi, yi), radius=0.04,
                            color=BLUE_C, fill_opacity=0.6) for xi, yi in zip(X, y)])
        self.add(dots)

        def make_line(w, b, c, wd):
            return Line(axes.coords_to_point(x_min, w * x_min + b),
                        axes.coords_to_point(x_max, w * x_max + b),
                        color=c, stroke_width=wd)

        self.add(make_line(TRUE_W, TRUE_B, RED, 3))
        fit_line = make_line(ws[0], bs[0], YELLOW, 3)
        self.add(fit_line)

        # 右上角显示 w, b, loss
        def make_info(i):
            return VGroup(
                Text(f"w = {ws[i]:.2f}, b = {bs[i]:.2f}", font_size=24, color=YELLOW),
                Text(f"MSE = {losses[i]:.2f}", font_size=24, color=YELLOW),
            ).arrange(DOWN, aligned_edge=LEFT, buff=0.1).to_corner(UR, buff=0.3)

        info = make_info(0)
        self.add(info)

        n = len(ws)
        tracker = ValueTracker(0)

        def upd_line(mob):
            i = min(int(tracker.get_value()), n - 1)
            mob.become(make_line(ws[i], bs[i], YELLOW, 3))

        def upd_info(mob):
            mob.become(make_info(min(int(tracker.get_value()), n - 1)))

        self.play(
            UpdateFromFunc(fit_line, upd_line),
            UpdateFromFunc(info, upd_info),
            tracker.animate.set_value(n - 1),
            run_time=6, rate_func=smooth,
        )
        self.wait(1)

# 对比：高斯噪声 vs 拉普拉斯噪声下，MSE 和 MAE 的表现
def gradient_descent_mae(X, y, w_init, b_init, lr, epochs):
    """使用 MAE 损失的梯度下降（次梯度）"""
    w, b = w_init, b_init
    w_hist, b_hist, loss_hist = [w], [b], [np.mean(np.abs(w * X + b - y))]
    
    for _ in range(epochs):
        y_pred = w * X + b
        error = y_pred - y
        # MAE 的次梯度：sign(error)
        dw = np.mean(np.sign(error) * X)
        db = np.mean(np.sign(error))
        w -= lr * dw
        b -= lr * db
        w_hist.append(w); b_hist.append(b)
        loss_hist.append(np.mean(np.abs(error)))
    return np.array(w_hist), np.array(b_hist), np.array(loss_hist)


# 生成拉普拉斯噪声数据（scale 比高斯 σ 大，且手动加几个极端离群点）
np.random.seed(123)
noise_laplace = np.random.laplace(0, NOISE_STD * 1.5, N)  # 拉普拉斯 scale 更大
y_laplace = TRUE_W * X + TRUE_B + noise_laplace

# 手动注入几个极端离群点——模拟「急用钱低价抛售」的异常房源
outlier_idx = [10, 45, 78]
y_laplace[outlier_idx] = TRUE_W * X[outlier_idx] + TRUE_B + np.array([25, -20, 18])

# 用 MSE 梯度下降拟合拉普拉斯噪声数据
w_mse, b_mse, _ = gradient_descent(X, y_laplace, 0.0, 0.0, 0.01, 1000)
# 用 MAE 梯度下降拟合拉普拉斯噪声数据（MAE 收敛慢，多跑几步）
w_mae, b_mae, _ = gradient_descent_mae(X, y_laplace, 0.0, 0.0, 0.005, 2000)

print(f"\n真实参数: w={TRUE_W}, b={TRUE_B}")
print(f"MSE 拟合:  w={w_mse[-1]:.3f}, b={b_mse[-1]:.3f}  (被离群点拉偏了)")
print(f"MAE 拟合:  w={w_mae[-1]:.3f}, b={b_mae[-1]:.3f}  (更接近真实值)")

# 可视化对比
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# 左图：高斯噪声数据（MSE 最优）
ax = axes[0]
ax.scatter(X, y, alpha=0.5, label='数据（高斯噪声）')
ax.plot(x_line, TRUE_W * x_line + TRUE_B, 'r--', lw=2, label='真实直线')
ax.plot(x_line, w_hist[-1] * x_line + b_hist[-1], 'g-', lw=2, label=f'MSE 拟合: y={w_hist[-1]:.2f}x+{b_hist[-1]:.2f}')
ax.legend(); ax.set_title('高斯噪声 → MSE 拟合')
ax.grid(True, alpha=0.3)

# 右图：拉普拉斯噪声 + 离群点 → MAE 更鲁棒
ax = axes[1]
ax.scatter(X, y_laplace, alpha=0.5, label='数据（含离群点）', c='gray')
# 高亮离群点
ax.scatter(X[outlier_idx], y_laplace[outlier_idx], alpha=1.0, s=100,
           label='离群点', c='red', marker='x', linewidths=2)
ax.plot(x_line, TRUE_W * x_line + TRUE_B, 'r--', lw=2, label='真实直线')
ax.plot(x_line, w_mse[-1] * x_line + b_mse[-1], 'g-', lw=2, label=f'MSE: y={w_mse[-1]:.2f}x+{b_mse[-1]:.2f}')
ax.plot(x_line, w_mae[-1] * x_line + b_mae[-1], 'orange', lw=2, label=f'MAE: y={w_mae[-1]:.2f}x+{b_mae[-1]:.2f}')
ax.legend(); ax.set_title('拉普拉斯噪声 + 离群点 → MAE 更鲁棒')
ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

# 最小二乘几何：N=3个数据点，D=2个参数（w, b）
# 在 R^3 中，Phi 的列空间是一个过原点的平面
# y 在这个平面上的投影就是最小二乘解

from mpl_toolkits.mplot3d import Axes3D

# 取前 3 个数据点做可视化
X3d = X[:3]
y3d = y[:3]

# 设计矩阵 [1, x]
Phi_3d = np.column_stack([np.ones(3), X3d])

# 正规方程求解
w_ols = np.linalg.solve(Phi_3d.T @ Phi_3d, Phi_3d.T @ y3d)
y_proj = Phi_3d @ w_ols

print(f"OLS 解: w={w_ols[1]:.3f}, b={w_ols[0]:.3f}")
print(f"y        = {y3d}")
print(f"投影 y_hat = {y_proj}")
print(f"残差      = {y3d - y_proj}")
print(f"残差 ⋅ 列1 = {np.dot(y3d - y_proj, Phi_3d[:, 0]):.2e}  (≈0 ✓)")
print(f"残差 ⋅ 列2 = {np.dot(y3d - y_proj, Phi_3d[:, 1]):.2e}  (≈0 ✓)")

# 3D 可视化
fig = plt.figure(figsize=(12, 10))
ax = fig.add_subplot(111, projection='3d')

# 列空间平面: span{(1,1,1), (x1,x2,x3)} 由 Phi 的两列张成
# 参数化: u*[1,1,1] + v*[x1,x2,x3]
u = np.linspace(-5, 5, 20)
v = np.linspace(-5, 5, 20)
U, V = np.meshgrid(u, v)
plane_x = U * Phi_3d[0, 0] + V * Phi_3d[0, 1]
plane_y = U * Phi_3d[1, 0] + V * Phi_3d[1, 1]
plane_z = U * Phi_3d[2, 0] + V * Phi_3d[2, 1]
ax.plot_surface(plane_x, plane_y, plane_z, alpha=0.2, color='blue')

# 原点
ax.scatter([0], [0], [0], color='black', s=80, marker='o')
ax.text(0, 0, 0, '  O', fontsize=12)

# y 向量
ax.quiver(0, 0, 0, y3d[0], y3d[1], y3d[2],
          color='red', linewidth=2, arrow_length_ratio=0.1, label='y')
ax.text(y3d[0], y3d[1], y3d[2], '  y', color='red', fontsize=14)

# 投影 y_hat
ax.quiver(0, 0, 0, y_proj[0], y_proj[1], y_proj[2],
          color='green', linewidth=2, arrow_length_ratio=0.1, label='ŷ = Proj(y)')
ax.text(y_proj[0], y_proj[1], y_proj[2], '  ŷ', color='green', fontsize=14)

# 残差（从 y_hat 到 y）
ax.quiver(y_proj[0], y_proj[1], y_proj[2],
          y3d[0] - y_proj[0], y3d[1] - y_proj[1], y3d[2] - y_proj[2],
          color='orange', linewidth=2, arrow_length_ratio=0.1, label='残差 (⊥ 列空间)')

# 列空间的两个基向量
col1 = Phi_3d[:, 0] * 3  # 缩放一下方便看
col2 = Phi_3d[:, 1] * 3
ax.quiver(0, 0, 0, col1[0], col1[1], col1[2],
          color='blue', linewidth=1.5, alpha=0.6, arrow_length_ratio=0.1)
ax.quiver(0, 0, 0, col2[0], col2[1], col2[2],
          color='blue', linewidth=1.5, alpha=0.6, arrow_length_ratio=0.1)

ax.set_xlabel('维度 1'); ax.set_ylabel('维度 2'); ax.set_zlabel('维度 3')
ax.set_title('最小二乘的几何：ŷ 是 y 在 Φ 列空间上的正交投影')
ax.legend()
plt.show()

def sgd(X, y, w_init, b_init, lr, epochs):
    """随机梯度下降：每步只看一个数据点"""
    w, b = w_init, b_init
    N = len(X)
    w_hist, b_hist, loss_hist = [w], [b], [compute_loss(X, y, w, b)]
    
    for epoch in range(epochs):
        i = np.random.randint(0, N)  # 随机挑一个点
        xi, yi = X[i], y[i]
        error = w * xi + b - yi
        w -= lr * 2 * error * xi
        b -= lr * 2 * error
        
        w_hist.append(w)
        b_hist.append(b)
        loss_hist.append(compute_loss(X, y, w, b))
    
    return np.array(w_hist), np.array(b_hist), np.array(loss_hist)


# 对比：批量梯度下降 vs SGD
w_batch, b_batch, loss_batch = gradient_descent(X, y, 0.0, 0.0, lr=0.01, epochs=500)
w_sgd, b_sgd, loss_sgd = sgd(X, y, 0.0, 0.0, lr=0.01, epochs=500)

fig, axes = plt.subplots(1, 3, figsize=(18, 5))

# Loss 曲线
ax = axes[0]
ax.plot(loss_batch, 'b-', alpha=0.7, linewidth=1.5, label='批量 GD')
ax.plot(loss_sgd, 'r-', alpha=0.5, linewidth=0.8, label='SGD')
ax.set_yscale('log')
ax.set_xlabel('Epoch'); ax.set_ylabel('MSE Loss')
ax.set_title('Loss 收敛对比')
ax.legend(); ax.grid(True, alpha=0.3)

# w 轨迹
ax = axes[1]
ax.plot(w_batch, 'b-', alpha=0.7, label='批量 GD')
ax.plot(w_sgd, 'r-', alpha=0.5, label='SGD')
ax.axhline(y=TRUE_W, color='gray', linestyle='--', label=f'真实 w={TRUE_W}')
ax.set_xlabel('Epoch'); ax.set_ylabel('w')
ax.set_title('w 的收敛轨迹（SGD 更吵但方向对）')
ax.legend(); ax.grid(True, alpha=0.3)

# b 轨迹
ax = axes[2]
ax.plot(b_batch, 'b-', alpha=0.7, label='批量 GD')
ax.plot(b_sgd, 'r-', alpha=0.5, label='SGD')
ax.axhline(y=TRUE_B, color='gray', linestyle='--', label=f'真实 b={TRUE_B}')
ax.set_xlabel('Epoch'); ax.set_ylabel('b')
ax.set_title('b 的收敛轨迹')
ax.legend(); ax.grid(True, alpha=0.3)

plt.suptitle('批量梯度下降 vs 随机梯度下降（SGD 吵但省计算）', fontsize=14, y=1.02)
plt.tight_layout()
plt.show()

print(f"批量 GD 最终: w={w_batch[-1]:.4f}, b={b_batch[-1]:.4f}")
print(f"SGD 最终:     w={w_sgd[-1]:.4f}, b={b_sgd[-1]:.4f}")
print(f"真实值:       w={TRUE_W}, b={TRUE_B}")

# 生成非线性数据：y = sin(x) + noise
N2 = 100
X2 = np.random.uniform(-3, 3, N2)
y2 = np.sin(X2) + np.random.normal(0, 0.15, N2)

plt.figure(figsize=(8, 5))
x_smooth = np.linspace(-3.5, 3.5, 200)
plt.plot(x_smooth, np.sin(x_smooth), 'r--', linewidth=2, label='真实函数: y = sin(x)')
plt.scatter(X2, y2, alpha=0.6, label='带噪声的数据点')
plt.xlabel('x'); plt.ylabel('y')
plt.legend(); plt.title('第 2 周：非线性数据')
plt.grid(True, alpha=0.3)
plt.show()

def poly_features(X, degree):
    """将 X 展开为多项式特征矩阵 [1, x, x^2, ..., x^degree]"""
    return np.vstack([X ** d for d in range(degree + 1)]).T  # shape: (N, degree+1)


def standardize(Phi):
    """对设计矩阵做标准化（每列减均值除以标准差），第一列（偏置列）不处理"""
    mean = np.mean(Phi[:, 1:], axis=0)
    std = np.std(Phi[:, 1:], axis=0)
    std[std == 0] = 1.0  # 防止除以 0
    Phi_scaled = Phi.copy()
    Phi_scaled[:, 1:] = (Phi[:, 1:] - mean) / std
    return Phi_scaled, mean, std


def poly_gradient_descent(Phi, y, lr, epochs, standardize_features=True, verbose=True):
    """对设计矩阵 Phi 做线性回归的梯度下降（可选标准化）"""
    if standardize_features and Phi.shape[1] > 1:
        Phi, mean, std = standardize(Phi)
    else:
        mean, std = None, None
    
    N, D = Phi.shape
    w = np.zeros(D)
    loss_history = [np.mean((Phi @ w - y) ** 2)]

    for epoch in range(epochs):
        error = Phi @ w - y
        grad = (2 / N) * Phi.T @ error
        w -= lr * grad
        loss_history.append(np.mean((Phi @ w - y) ** 2))

        if verbose and epoch % max(1, epochs // 10) == 0:
            print(f"  Epoch {epoch:4d}: loss={loss_history[-1]:.6f}")

    return w, np.array(loss_history), mean, std


def poly_predict(X, w, mean=None, std=None):
    """用多项式权重做预测（支持反标准化）"""
    degree = len(w) - 1
    Phi = poly_features(X, degree)
    if mean is not None and std is not None:
        Phi_scaled = Phi.copy()
        Phi_scaled[:, 1:] = (Phi[:, 1:] - mean) / std
        return Phi_scaled @ w
    return Phi @ w

# 用不同多项式次数拟合，观察过拟合
degrees = [1, 3, 5, 9, 15]
lr_poly = 0.1
epochs_poly = 2000

fig, axes = plt.subplots(2, 3, figsize=(16, 10))
axes = axes.flatten()

results = {}

for idx, d in enumerate(degrees):
    Phi = poly_features(X2, d)
    w, loss_hist, mean, std = poly_gradient_descent(Phi, y2, lr_poly, epochs_poly)

    # 预测
    X_test = np.linspace(-3.5, 3.5, 200)
    y_pred = poly_predict(X_test, w, mean, std)

    # 训练集上的预测（用标准化后的 Phi）
    if d > 1:
        Phi_scaled = Phi.copy()
        Phi_scaled[:, 1:] = (Phi[:, 1:] - mean) / std
        y_train_pred = Phi_scaled @ w
    else:
        y_train_pred = Phi @ w
    train_loss = np.mean((y_train_pred - y2) ** 2)

    # 画拟合曲线
    ax = axes[idx]
    ax.scatter(X2, y2, alpha=0.5, s=20, label='训练数据')
    ax.plot(x_smooth, np.sin(x_smooth), 'r--', lw=2, label='真实 sin(x)')
    ax.plot(X_test, y_pred, 'g-', lw=2, label=f'd={d} 拟合')
    ax.set_xlim(-3.5, 3.5); ax.set_ylim(-1.8, 1.8)
    ax.set_title(f'次数 d = {d}  |  Train Loss = {train_loss:.4f}')
    ax.legend(fontsize=8)
    ax.grid(True, alpha=0.3)

    results[d] = {'w': w, 'loss': train_loss, 'loss_hist': loss_hist, 'mean': mean, 'std': std}

# 最后一个子图：各次数的训练损失对比
ax = axes[-1]
for d in degrees:
    ax.plot(results[d]['loss_hist'], label=f'd={d}', alpha=0.8)
ax.set_xlabel('Epoch'); ax.set_ylabel('MSE Loss')
ax.set_title('各次数训练损失曲线')
ax.set_yscale('log')
ax.legend(fontsize=8)
ax.grid(True, alpha=0.3)

plt.suptitle('多项式拟合：从欠拟合到过拟合', fontsize=18, y=1.02)
plt.tight_layout()
plt.show()

# 可视化各次数下权重的范数
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# 左图：权重 L2 范数
ax = axes[0]
l2_norms = [np.linalg.norm(results[d]['w']) for d in degrees]
bars = ax.bar(degrees, l2_norms, color=['steelblue', 'steelblue', 'steelblue', 'coral', 'firebrick'])
ax.set_xlabel('多项式次数 d'); ax.set_ylabel('||w||₂')
ax.set_title('权重 L2 范数随多项式次数增长')
for bar, val in zip(bars, l2_norms):
    ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.3, f'{val:.1f}', ha='center', fontsize=10)

# 右图：各次数权重的具体值
ax = axes[1]
for d in degrees:
    w = results[d]['w']
    ax.plot(range(len(w)), w, 'o-', markersize=6, label=f'd={d}', alpha=0.8)
ax.axhline(y=0, color='gray', linestyle='-', alpha=0.3)
ax.set_xlabel('权重索引'); ax.set_ylabel('权重值')
ax.set_title('各次数下权重分布')
ax.legend(fontsize=9)
ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

print("结论：d 越大，||w||₂ 越大 → 模型越容易对训练数据「死记硬背」 → 过拟合")
print("解决思路：在损失函数中加罚项 λ||w||² → 这就是 L2 正则化（Ridge）")

def ridge_gradient_descent(Phi, y, lr, epochs, lam, standardize_features=True):
    """带 L2 正则化的梯度下降（Ridge），可选标准化"""
    if standardize_features and Phi.shape[1] > 1:
        Phi, mean, std = standardize(Phi)
    else:
        mean, std = None, None
    
    N, D = Phi.shape
    w = np.zeros(D)
    loss_history = []

    for epoch in range(epochs):
        error = Phi @ w - y
        # 梯度 = MSE 梯度 + λ * 2w（L2 惩罚的梯度）
        grad_mse = (2 / N) * Phi.T @ error
        grad_l2 = 2 * lam * w
        grad = grad_mse + grad_l2
        w -= lr * grad

        mse = np.mean(error ** 2)
        l2_penalty = lam * np.sum(w ** 2)
        loss_history.append(mse + l2_penalty)

    return w, np.array(loss_history), mean, std


# 用 d=15 的高次多项式，对比有无 L2 正则化
d_high = 15
Phi_high = poly_features(X2, d_high)

# 无正则化（用标准化）
w_no_reg, loss_no_reg, mean_ridge, std_ridge = poly_gradient_descent(Phi_high, y2, lr=0.1, epochs=3000, standardize_features=True)

# 有正则化（不同 λ）
lambdas = [0.0, 0.001, 0.01, 0.1, 1.0]
ridge_results = {}
for lam in lambdas:
    if lam == 0.0:
        w_r, loss_r = w_no_reg, loss_no_reg
    else:
        w_r, loss_r, _, _ = ridge_gradient_descent(Phi_high, y2, lr=0.1, epochs=3000, lam=lam, standardize_features=True)
    ridge_results[lam] = {'w': w_r, 'loss': loss_r}

# 展示：正则化如何在梯度更新中压制大权重
# 取 d=15 训练的第一步，对比 λ=0 和 λ=1.0 的梯度
# 标准化 Phi_high（和前面的 Ridge 训练保持一致）
Phi_s = Phi_high.copy()
Phi_s[:, 1:] = (Phi_s[:, 1:] - mean_ridge) / std_ridge
w0 = np.zeros(Phi_s.shape[1])

# MSE 梯度（不含正则化）
err = Phi_s @ w0 - y2
grad_mse = (2 / len(y2)) * Phi_s.T @ err

# L2 惩罚梯度 = 2λw，初始 w=0 时惩罚梯度也是 0
# 但跑几步之后 w 变大了，惩罚梯度就开始起作用
# 我们模拟跑 10 步后的状态
w_10 = w0.copy()
for _ in range(10):
    err = Phi_s @ w_10 - y2
    w_10 -= 0.1 * ((2/len(y2)) * Phi_s.T @ err)  # 只用 MSE

# 此时分别计算 MSE 梯度和 L2 惩罚梯度
err_10 = Phi_s @ w_10 - y2
grad_mse_10 = (2 / len(y2)) * Phi_s.T @ err_10
grad_l2_10 = 2 * 1.0 * w_10  # λ=1.0 时的惩罚梯度

# 画前 10 个权重的梯度对比
fig, ax = plt.subplots(figsize=(10, 5))
x = np.arange(10)
w = 0.35
ax.bar(x - w/2, np.abs(grad_mse_10[:10]), w, label='MSE 梯度（数据要你改多少）', color='steelblue')
ax.bar(x + w/2, np.abs(grad_l2_10[:10]), w, label='L2 惩罚梯度（λ=1.0 拉回零的力）', color='coral')
ax.set_xlabel('权重索引'); ax.set_ylabel('梯度绝对值')
ax.set_title('梯度分解：MSE 想放大权重，L2 把它们拽回来')
ax.legend(); ax.grid(True, alpha=0.3)
plt.show()

print("MSE 梯度让权重变大（拟合数据），L2 惩罚梯度反向拉回（防止过拟合）。")
print("权重越大 → L2 惩罚梯度越大 → 被拉回越狠。这就是正则化的机制。")

# 可视化：不同 λ 下 d=15 多项式的拟合效果
X_test = np.linspace(-3.5, 3.5, 200)
Phi_test = poly_features(X_test, d_high)
# 标准化测试特征
Phi_test_scaled = Phi_test.copy()
Phi_test_scaled[:, 1:] = (Phi_test[:, 1:] - mean_ridge) / std_ridge
# 标准化训练特征
Phi_high_scaled = Phi_high.copy()
Phi_high_scaled[:, 1:] = (Phi_high[:, 1:] - mean_ridge) / std_ridge

fig, axes = plt.subplots(2, 3, figsize=(16, 10))
axes = axes.flatten()

for idx, lam in enumerate(lambdas):
    w = ridge_results[lam]['w']
    y_pred = Phi_test_scaled @ w
    train_pred = Phi_high_scaled @ w
    train_mse = np.mean((train_pred - y2) ** 2)

    ax = axes[idx]
    ax.scatter(X2, y2, alpha=0.5, s=20, label='训练数据')
    ax.plot(x_smooth, np.sin(x_smooth), 'r--', lw=2, label='真实 sin(x)')
    ax.plot(X_test, y_pred, 'g-', lw=2, label=f'λ={lam}')
    ax.set_xlim(-3.5, 3.5); ax.set_ylim(-1.8, 1.8)
    title = '无正则化 (过拟合!)' if lam == 0.0 else f'Ridge λ={lam}'
    ax.set_title(f'{title}\n||w||₂ = {np.linalg.norm(w):.2f}, Train MSE = {train_mse:.4f}')
    ax.legend(fontsize=8)
    ax.grid(True, alpha=0.3)

# 最后一个子图：各 λ 下的权重分布
ax = axes[-1]
for lam in lambdas:
    w = ridge_results[lam]['w']
    ax.plot(range(len(w)), w, 'o-', markersize=4, label=f'λ={lam}', alpha=0.8)
ax.axhline(y=0, color='gray', linestyle='-', alpha=0.3)
ax.set_xlabel('权重索引'); ax.set_ylabel('权重值')
ax.set_title('不同 λ 下的权重分布（d=15）')
ax.legend(fontsize=8)
ax.grid(True, alpha=0.3)

plt.suptitle('L2 正则化效果：λ 越大 → 权重越小 → 曲线越平滑', fontsize=16, y=1.02)
plt.tight_layout()
plt.show()

# L1 vs L2 几何可视化：约束区域 + MSE 等高线
# 在二维参数空间中展示 Ridge (圆) 和 Lasso (菱形) 的区别

# 用 sin 数据的 d=1 线性模型（只有 w 和 b 两个参数），画 MSE 等高线
# 然后叠加 L1/L2 约束区域

# 在 (w, b) 网格上计算 MSE
w_grid = np.linspace(-2, 6, 200)
b_grid = np.linspace(-2, 8, 200)
W, B = np.meshgrid(w_grid, b_grid)
MSE_grid = np.zeros_like(W)
for i in range(len(w_grid)):
    for j in range(len(b_grid)):
        MSE_grid[j, i] = np.mean((W[j, i] * X2 + B[j, i] - y2) ** 2)

# MSE 最小值（无正则化时的最优解）
# 用正规方程求解
Phi_simple = np.column_stack([X2, np.ones_like(X2)])  # [x, 1]，注意顺序：w 在 x 列，b 在 1 列
w_ols_simple = np.linalg.solve(Phi_simple.T @ Phi_simple, Phi_simple.T @ y2)
w_opt, b_opt = w_ols_simple[0], w_ols_simple[1]

fig, axes = plt.subplots(1, 2, figsize=(16, 7))

# 公共的 MSE 等高线级别
levels = np.geomspace(MSE_grid.min() + 0.01, MSE_grid.max() * 0.3, 12)

for ax, penalty, title, shape_label in [
    (axes[0], 'l2', 'Ridge (L2)：圆约束', '圆 ∥w∥₂ ≤ t'),
    (axes[1], 'l1', 'Lasso (L1)：菱形约束', '菱形 ∥w∥₁ ≤ t'),
]:
    # MSE 等高线
    cs = ax.contour(W, B, MSE_grid, levels=levels, colors='steelblue', alpha=0.6, linewidths=1)
    ax.clabel(cs, inline=True, fontsize=7, fmt='%.1f')
    
    # 无约束最优解
    ax.plot(w_opt, b_opt, 'k*', markersize=15, label=f'无约束最优 ({w_opt:.2f}, {b_opt:.2f})')
    
    # 约束区域（三个不同 t）
    ts = [1.5, 3.0, 5.0]
    colors = ['#ff6b6b', '#ffa500', '#2ecc71']
    for t, c in zip(ts, colors):
        theta = np.linspace(0, 2*np.pi, 300)
        if penalty == 'l2':
            # L2: w² + b² ≤ t² 但注意我们中心在 (0, 0)
            x_bound = w_opt + t * np.cos(theta)
            y_bound = b_opt + t * np.sin(theta)
        else:
            # L1: |w-w_opt| + |b-b_opt| ≤ t → 菱形
            # 菱形顶点: right, top, left, bottom
            vertices = np.array([
                [w_opt + t, b_opt],
                [w_opt, b_opt + t],
                [w_opt - t, b_opt],
                [w_opt, b_opt - t],
                [w_opt + t, b_opt],
            ])
            x_bound, y_bound = vertices[:, 0], vertices[:, 1]
        
        ax.fill(x_bound, y_bound, alpha=0.12, color=c, edgecolor=c, linewidth=2,
                label=f't={t}')
    
    ax.set_xlabel('w'); ax.set_ylabel('b')
    ax.set_title(title)
    ax.legend(fontsize=9)
    ax.grid(True, alpha=0.3)
    ax.set_xlim(w_opt - 7, w_opt + 7)
    ax.set_ylim(b_opt - 7, b_opt + 7)

plt.suptitle('L1 vs L2 正则化的几何直觉：MSE 等高线 + 约束区域', fontsize=16, y=1.02)
plt.tight_layout()
plt.show()

print("关键观察：")
print(f"  无约束 MSE 最优解: w={w_opt:.2f}, b={b_opt:.2f}")
print(f"  Ridge (L2):  约束是圆，MSE 等高线碰到圆弧 → 两个参数都收缩，但都不为零")
print(f"  Lasso (L1):  约束是菱形，MSE 等高线碰到尖角 → 某些参数恰好为零")
print(f"\n  这在 PRML Fig 3.3 和 Fig 3.4 有详细说明。")

# 偏差-方差分解实验：从同一真实函数重复采样 → 训练 → 统计预测分布
def bias_variance_experiment(true_func, n_datasets=100, n_train=30, noise_std=0.15,
                              degrees=[1, 3, 5, 9, 15], x_range=(-3, 3)):
    """对每个多项式次数 d，在 n_datasets 组数据上训练，计算 bias² 和 variance"""
    np.random.seed(42)
    x_test = np.linspace(x_range[0], x_range[1], 200)
    
    results = {}
    for d in degrees:
        predictions = np.zeros((n_datasets, len(x_test)))
        
        for ds in range(n_datasets):
            # 采样训练集
            X_train = np.random.uniform(*x_range, n_train)
            y_train = true_func(X_train) + np.random.normal(0, noise_std, n_train)
            
            # 训练
            Phi_train = poly_features(X_train, d)
            w, _, mean, std = poly_gradient_descent(Phi_train, y_train,
                                                     lr=0.1, epochs=2000, standardize_features=(d > 1), verbose=False)
            # 预测
            y_pred = poly_predict(x_test, w, mean, std)
            predictions[ds] = y_pred
        
        # 统计
        mean_pred = predictions.mean(axis=0)
        bias_sq = (mean_pred - true_func(x_test)) ** 2
        variance = predictions.var(axis=0)
        avg_bias_sq = bias_sq.mean()
        avg_variance = variance.mean()
        
        results[d] = {
            'predictions': predictions,
            'mean_pred': mean_pred,
            'bias_sq': bias_sq,
            'variance': variance,
            'avg_bias_sq': avg_bias_sq,
            'avg_variance': avg_variance,
        }
    
    return results, x_test


# 跑实验
true_func = np.sin
bv_results, x_bv = bias_variance_experiment(true_func, n_datasets=100, n_train=30)
degrees_bv = list(bv_results.keys())

# 打印各 d 的 bias² 和 variance
print(f"{'d':>3}  {'bias²':>10}  {'variance':>10}  {'bias²+variance':>15}")
print("-" * 42)
for d in degrees_bv:
    b2 = bv_results[d]['avg_bias_sq']
    v = bv_results[d]['avg_variance']
    print(f"{d:>3}  {b2:>10.5f}  {v:>10.5f}  {b2+v:>15.5f}")