# 生成数据
import torch
from matplotlib import pyplot as plt
import random
import traceback

# create data
def create_data(W, b, num):
    X = torch.normal(mean=0, std=1, size =(num, len(W)))
    y = X.matmul(W) + b
    # 加点噪声
    y += torch.normal(mean=0, std=0.1, size=(num,))
    return X, y 

def plot_scatt(x, y):
    plt.scatter(x, y)
                      
W_true = torch.tensor([3, -20.5])
b_true = 8
data_num = 1000
features, labels = create_data(W_true, b_true, data_num)
plot_scatt(features[:, 1], labels)
plot_scatt(features[:, 0], labels)

# **** 优化1，使用api读取batch数据
from torch.utils import data
def get_batch(data_arrays, batch_size, is_train=True):
    data_set = data.TensorDataset(*data_arrays)
    return data.DataLoader(data_set, batch_size, shuffle=is_train)
        
batch_size = 10
data_iter = get_batch((features, labels), batch_size)
print(type(data_iter), next(iter(data_iter)))

from torch import nn


# model
net = nn.Sequential(nn.Linear(2, 1))
# net[0]指网络第一层
net[0].weight.data.normal_(0, 0.01)
net[0].bias.data.fill_(0)

# loss
loss = nn.MSELoss()

# sgd
trainer = torch.optim.SGD(net.parameters(), lr=0.1)
    

# 训练过程
num_epochs = 30
for epoch in range(num_epochs):
    for X, y in data_iter:
        l = loss(net(X), y)
        trainer.zero_grad()
        l.backward()
        trainer.step()
    l = loss(net(features), labels)
    print(f'epoch {epoch + 1}, loss {l:f}')

W = net[0].weight.data
b = net[0].bias.data
print(f'w的估计误差: {W_true - W.reshape(W_true.shape)}')
print(f'b的估计误差: {b_true - b}')

# def plot_linear(x, y, w, b):
#     plt.scatter(x, y)
#     yy_list = []
#     for xx in x:
#         yy_list.append(w * xx.item() +b.item())
# #     print(yy_list)
#     plt.plot(x, yy_list)
    
# # print(W[1].item())
# plot_linear(features[:, 1], labels, W[1].item(), b)
# plot_linear(features[:, 0], labels, W[0].item(), b)