import numpy as np
import matplotlib.pyplot as plt

# 步骤1: 生成模拟数据
np.random.seed(0)
X = 2 * np.random.rand(100, 1)
y = 4 + 3 * X + 2 * X**2 + np.random.randn(100, 1)

# 步骤2: 定义线性模型 (我们从随机权重开始)
w = np.random.randn(2, 1)
b = np.random.randn(1)


# 步骤3: 定义损失函数 (均方误差)
def mse(y_pred, y_true):
    return ((y_pred - y_true) ** 2).mean()


# 步骤4: 使用梯度下降来优化模型参数
learning_rate = 0.01
epochs = 100

# 记录每次迭代的损失，用于画图
losses = []

# 准备一个颜色列表，用于绘制不同颜色的直线
colors = plt.cm.rainbow(np.linspace(0, 1, epochs))

fig, ax = plt.subplots()

# 绘制数据点
ax.scatter(X, y, label="Data")

# 扩展X值以绘制连续的线条
X_plot = np.linspace(X.min(), X.max(), 100).reshape(-1, 1)
X_plot_ext = np.hstack((X_plot, X_plot**2))

for epoch in range(epochs):
    # 扩展特征
    X_ext = np.hstack((X, X**2))

    # 预测
    y_pred = X_ext @ w + b

    # 计算损失
    loss = mse(y_pred, y)
    losses.append(loss)

    # 梯度下降
    dy_pred = 2 * (y_pred - y)
    dw = X_ext.T @ dy_pred / len(X_ext)
    db = dy_pred.sum(axis=0) / len(X_ext)

    # 更新权重和偏置
    w -= learning_rate * dw
    b -= learning_rate * db

    # 绘制当前拟合的直线，使用不同的颜色
    y_plot_pred = X_plot_ext @ w + b
    ax.plot(X_plot, y_plot_pred, color=colors[epoch], label=f"Epoch {epoch + 1}")

# 设置图例和标签
ax.set_xlabel("X")
ax.set_ylabel("y")
ax.legend()

plt.show()

# 绘制损失下降的曲线
plt.plot(losses)
plt.title("Loss over epochs")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.show()