import torch
import matplotlib.pyplot as plt #  画图的

import random  #随机

# 要预测的值是w，就是权重的值
def create_data(w, b, data_num):  #生成数据
    # 生成数据，均值为0，方差为1，维度为数据数量，宽度为w的维度
    x = torch.normal(0, 1, (data_num, len(w)))

    # 矩阵相乘，在加上b
    y = torch.matmul(x, w) + b    #matmul表示矩阵相乘

    # 均值为0，方差为1的噪声，维度为y的维度
    noise = torch.normal(0, 0.01, y.shape)  #噪声要加到y上
    y += noise

    return x, y

num = 500

true_w = torch.tensor([8.1, 2, 2, 4])
true_b = torch.tensor(1.1)

# 表示数据库的自变量和标签
X, Y = create_data(true_w, true_b, num)

# 画一个散点图，点的大小为1，Y是标签，X目前为全部的行，第三列
# 不同的列对应不同的散点图，权值大的呈近似线性
# plt.scatter(X[:, 3], Y, 1)
# plt.show()
# 到目前为止，已经生成好了数据，并且自己设定好了标签值
# 下面的神经网络就这些数据来推测真实的上面自己设定的权值，看看准不准！！！


# 不再整体500个数据一起运算，每次取一部分进行预算
def data_provider(data, label, batchsize):       # 每次访问这个函数， 就能提供一批数据

    length = len(label)
    indices = list(range(length))

    # 我不能按顺序取  把数据打乱
    random.shuffle(indices)

    # 每次取出batchsize个数据和标签
    #
    for each in range(0, length, batchsize):
        get_indices = indices[each: each+batchsize]
        get_data = data[get_indices]
        get_label = label[get_indices]

        yield get_data, get_label  # 有存档点的return，功能就是return增加了存档点，下次访问函数循环，从yield开始，显示第二批数据

batchsize = 16
# for batch_x, batch_y in data_provider(X, Y, batchsize):
#     print(batch_x, batch_y)
#     break

# 预测y = wx + b
def fun(x, w, b):
    pred_y = torch.matmul(x, w) + b
    return pred_y

# 计算loss，得到平均loss，接下来就进行梯度下降
def maeLoss(pre_y, y):
    return torch.sum(abs(pre_y - y)) / len(y)

# 梯度下降
def sgd(paras, lr): # 随机梯度下降，更新参数
    with torch.no_grad():  # 属于这句代码的部分，不计算梯度
        # 分别对每个参数减去自己的梯度乘以学习率，就是梯度下降
        for para in paras: # paras为[w, b]
            para -= para.grad * lr      # 不能写成   para = para - para.grad*lr
            para.grad.zero_()      # 使用过的梯度，归0，否则会累加

# 学习率
lr = 0.03

# 这是随机的值，用来进行后续的模型调优
w_0 = torch.normal(0, 0.01, true_w.shape, requires_grad=True)   # 这个w需要计算梯度
b_0 = torch.tensor(0.01, requires_grad=True)
# print(w_0, b_0)

# 要训练的次数
epochs = 50

# 开始训练
for epoch in range(epochs):
    data_loss = 0  # 每一批数据计算一次loss总值
    for batch_x, batch_y in data_provider(X, Y, batchsize):  # batch_x，batch_y为数据集的数据和标签
        pred_y = fun(batch_x,w_0, b_0)  # 算出预测值
        loss = maeLoss(pred_y, batch_y)  # 计算差值
        loss.backward()  # 梯度回传，所有的参数的梯度值都会被算出来被存储到参数里，梯度值就是导数值
        sgd([w_0, b_0], lr)  # 更新模型
        # print(loss)
        data_loss += loss

    # print("epoch %03d: loss: %.6f" % (epoch, data_loss))
#
#
print("真实的函数值是", true_w, true_b)
print("训练得到的参数值是", w_0, b_0)

# 可视化
idx = 3
plt.plot(X[:, idx].detach().numpy(), X[:, idx].detach().numpy()*w_0[idx].detach().numpy()+b_0.detach().numpy())
plt.scatter(X[:, idx], Y, 1)
plt.show()