#多项式拟合
import torch
import numpy as np

from FC_ML_Data.FC_ML_Data_Output.Data_Output_File import tensor_to_excel

# 真实多项式系数
true_w = torch.tensor([0.5, 3.0, 2.4])  # 对应x, x², x³项
true_b = 0.9

# 生成训练数据
def make_features(x):
    return torch.stack([x**i for i in range(1,4)], dim=1)  # 构建x, x², x³特征矩阵

x = torch.linspace(-3, 3, 100)
X = make_features(x)
y = X @ true_w + true_b + torch.randn(x.size()) * 0.5  # 添加噪声
print(x,X,y)
# tensor_to_excel(torch.cat([x, y], dim=-1),"./")


class PolyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(3, 1)  # 输入3维(x,x²,x³)，输出1维

    def forward(self, x):
        return self.linear(x)


model = PolyModel()
criterion = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

for epoch in range(1000):
    pred = model(X)
    loss = criterion(pred.squeeze(), y)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if epoch % 100 == 0:
        print(f'Epoch {epoch}, Loss: {loss.item():.4f}')

# 获取训练后的参数
w_pred = model.linear.weight.detach().numpy().flatten()
b_pred = model.linear.bias.detach().numpy()

# print(f"真实参数: w={true_w.numpy()}, b={true_b}")
# print(f"预测参数: w={w_pred}, b={b_pred:.2f}")

# 可视化
import matplotlib.pyplot as plt
plt.scatter(x, y, label='ori')
plt.plot(x, model(X).detach().numpy(), 'r', label='fit')
plt.legend()
plt.show()