import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
import torch
import random
import numpy

def create_data(data_num, true_w, true_b):
    X = torch.rand((data_num, len(true_w)))*49+1
    y = torch.matmul(X, true_w) + true_b
    noise = torch.normal(0, 2, y.shape)
    y += noise
    return X, y


data_num = 1000
true_w = torch.tensor([10.0, 6.0, -3.4, 7.0, 1.5])
true_b = torch.tensor(3.8)

X, y = create_data(data_num, true_w, true_b)
plt.scatter(X[:,0].detach().numpy(), y, 1)
plt.tight_layout()
plt.show()

# 提供一批数据
def data_provide(X, y, batch_size):
    indices = [i for i in range(len(X))]
    random.shuffle(indices)
    for each in range(0, len(X), batch_size):
        get_indices = indices[each : each+batch_size]
        get_X = X[get_indices]
        get_y = y[get_indices]
        yield get_X, get_y   # 下次再执行data_provide时直接跳到for each。each值继续

# 前向传播
def fun(X, w, b):
    return torch.matmul(X, w) + b

def loss(pred_y, y):
    return torch.sum(torch.abs(pred_y - y)) / len(y)

def sgd(paras, lr):
    with torch.no_grad():
        for para in paras:
            para -= para.grad*lr
            para.grad.zero_()


lr = 0.001
w0 = torch.randn(true_w.shape, requires_grad=True)
b0 = torch.rand(1, requires_grad=True)
batch_size = 25
epoches = 300

for i in range(epoches):
    loss_sum = 0
    for batch_X, batch_y in data_provide(X, y, batch_size):
        pred_y = fun(batch_X, w0, b0)
        loss_avg = loss(pred_y, batch_y)
        loss_avg.backward()
        sgd([w0, b0], lr)
        loss_sum += loss_avg
    print(f'第{i}轮，损失{loss_sum:.3f}.  w={w0},b={b0}')