import numpy as np
from torchvision.datasets import MNIST
from torchvision.transforms import Compose, ToTensor, Normalize
import torchvision
from torch.utils.data import DataLoader
import torch
import torch.nn.functional as F
from torch import nn
from torch.optim import Adam
import os

data = np.random.randint(0, 255, size=12)
img = data.reshape(2, 2, 3)
print(img.shape)
img_tensor = torchvision.transforms.ToTensor()(img)  # 转换成tensor类型
print(img_tensor)
print(img_tensor.shape)

mninst = MNIST(root='./data', train=True, download=True)
print(mninst[0])
ret = torchvision.transforms.ToTensor()(mninst[0][0])
print(ret.size())

data = np.random.randint(0, 255, size=12)
img = data.reshape(2, 2, 3)
img = torchvision.transforms.transforms.ToTensor()(img).float()  # 转成tensor
print(img)
normal_img = torchvision.transforms.Normalize((10, 10, 10), (1, 1, 1))(img)  # 进行规范化处理
print(normal_img)
b = torch.randn(2, 3)
print(b)
b = F.relu(b)
print(b)
print("-------------------------------开始实现手写数字的识别---------------------------------")
# 开始实现手写数字的识别

Batch_Size = 128


# 准备数据集，其中0.1307,0.3081为MNIST数据的均值和标准差，这样操作能够对其进行标准化
def get_dataloader(train=True):
    transform_fn = Compose([
        ToTensor(),
        Normalize(mean=(0.1307,), std=(0.3081,))  # mean和std的形状相同
    ])
    dataset = MNIST(root='/data', train=train, transform=transform_fn, download=True)
    data_loader = DataLoader(dataset, batch_size=Batch_Size, shuffle=True)
    return data_loader


data_loader = get_dataloader()


# 构建模型
class MnistModel(nn.Module):
    def __init__(self):
        super(MnistModel, self).__init__()
        self.fc1 = nn.Linear(1 * 28 * 28, 28)
        self.fc2 = nn.Linear(28, 10)

    def forward(self, input):  # input的形状为：[batch_size,1,28,28]
        # 1.形状的修改
        x = input.view([-1, 1 * 28 * 28])
        # 2.进行全连接操作
        x = self.fc1(x)
        # 3.进行激活函数的处理
        x = F.relu(x)  # 经过激活函数的处理，形状不会发生变化
        # 4.输出层
        out = self.fc2(x)
        return F.log_softmax(out, dim=-1)


model = MnistModel()
if os.path.exists("./model/mnist_model.pkl"):
    model.load_state_dict(torch.load("./model/mnist_model.pkl"))
optimizer = Adam(model.parameters(), lr=0.001)
if os.path.exists("./model/mnist_optimizer.pkl"):
    optimizer.load_state_dict(torch.load("./model/mnist_optimizer.pkl"))


# 模型的训练
def train(epoch):
    data_loader = get_dataloader()
    for index, (input, target) in enumerate(data_loader):
        output = model(input)  # 调用模型得到预测值
        loss = F.nll_loss(output, target)  # 得到损失函数
        optimizer.zero_grad()  # 将梯度置为0
        loss.backward()  # 反向传播
        optimizer.step()  # 梯度的更新
        if index % 100 == 0:
            print(epoch, index, loss.item())
        # 模型的保存
        if index % 100 == 0:  # 表示每隔100步保存一次
            torch.save(model.state_dict(), "./model/mnist_model.pkl")
            torch.save(optimizer.state_dict(), "./model/mnist_optimizer.pkl")


def test():
    loss_list = []
    acc_list = []
    test_dataloader = get_dataloader(train=False)
    for index, (input, target) in enumerate(test_dataloader):
        '''为了防止跟踪历史记录(和使用内存)，可以将代码块包装在with torch.no_grad():中。
        在评估模型时特别有用，因为模型可能具有requires_grad = True的可训练的参数，
        但是我们不需要在此过程中对他们进行梯度计算。 '''
        with torch.no_grad():
            output = model(input)  # output形状：[batchsize,10]
            cur_loss = F.nll_loss(output, target)  # target形状：[batchsize]
            loss_list.append(cur_loss)
            # 计算准确率
            predict = output.max(dim=-1)[-1]
            cur_acc = predict.eq(target).float().mean()
            acc_list.append(cur_acc)
    print("平均准确率：{}，平均损失：{}".format(np.mean(acc_list), np.mean(loss_list)))


if __name__ == '__main__':
    # for i in range(3):
    #     train(i)
    test()