import matplotlib.pyplot as plt

x_data = [1.0, 2.0, 3.0]
y_data = [2.0, 4.0, 6.0]

w = 1.0
#forward(x)：前向传播（模型预测）
def forward(x):
    return x * w  # 线性模型：y_pred = w * x
#loss(x, y)：损失函数，计算单个样本上的平方误差（MSE的组成部分）
def loss(x, y):
    y_pred = forward(x)
    return (y_pred - y) ** 2  # 单个样本损失： (预测值 - 真实值)^2
# gradient(x, y): 计算单个样本损失对参数 w 的梯度（导数）
def gradient(x, y):
    return 2 * x * (x * w - y)  # 公式推导：d(loss)/dw=2*x*(w*x-y)

# 记录绘图数据的列表
epoch_list = []     # 记录外层循环的epoch数
loss_list = []      # 记录每个epoch结束时的损失loss值
w_list = []         # 记录每个epoch结束时的w值

print('Predict (before training):', 4, forward(4))
for epoch in range(100): # 迭代100次
    for x, y in zip(x_data, y_data):
        l = loss(x, y) # 1. 计算当前损失
        grad = gradient(x, y) # 2. 计算当前梯度
        w -= 0.01 * grad  # 3. 核心：沿负梯度方向更新参数（# 0.01是学习率learning_rate，控制每一步更新的大小）
        print("/tgrad: ", x, y, grad)

    # ===== 记录每个epoch结束后的状态 =====
    epoch_list.append(epoch)   # 记录当前epoch
    loss_list.append(l)        # 记录最后一个样本的损失（近似epoch损失）
    w_list.append(w)           # 记录更新后的w
    print("progress: ", epoch, "w=", w, "loss=", l)
print('Predict (after training):', 4, forward(4))

# # ============= 绘图部分 =============
# 第一张图：损失下降
plt.figure(figsize=(12, 4))

# 左图：损失
plt.subplot(1, 2, 1)
plt.plot(epoch_list, loss_list, 'b-')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Loss')

# 右图：参数w变化
plt.subplot(1, 2, 2)
plt.plot(epoch_list, w_list, 'r-')
plt.xlabel('Epoch')
plt.ylabel('Weight w')
plt.title('Parameter w')

plt.show()