import torch
import torchvision
from torch import optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
from Pytorch_MNIST图片识别.model import Net
import matplotlib.pyplot as plt
import os

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

# TODO epoch的数量定义了我们将循环整个训练数据集的次数
n_epochs = 3

# TODO 使用batch_size=64进行训练，并使用size=1000对这个数据集进行测试
batch_size_train = 64
batch_size_test = 1000

# TODO 优化器的超参数
learning_rate = 0.01
momentum = 0.5

log_interval = 10
random_seed = 1
torch.manual_seed(random_seed)

# TODO 自动将MNIST数据集下载到目录下的data文件夹
train_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('./data/', train=True, download=True,
                               transform=torchvision.transforms.Compose([

                                   torchvision.transforms.ToTensor(),
                                   # TODO MNIST数据集的全局平均值和标准偏差
                                   torchvision.transforms.Normalize(
                                       (0.1307,), (0.3081,))
                               ])),
    batch_size=batch_size_train, shuffle=True)
test_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('./data/', train=False, download=True,
                               transform=torchvision.transforms.Compose([
                                   torchvision.transforms.ToTensor(),
                                   # TODO MNIST数据集的全局平均值和标准偏差
                                   torchvision.transforms.Normalize(
                                       (0.1307,), (0.3081,))
                               ])),
    batch_size=batch_size_test, shuffle=False)

# TODO 初始化网络和优化器
network = Net()
optimizer = optim.SGD(network.parameters(), lr=learning_rate, momentum=momentum)

train_losses = []
train_counter = []
test_losses = []
test_counter = [i * len(train_loader.dataset) for i in range(n_epochs + 1)]

# TODO 模型存储位置(一个是完整的模型，一个是只有参数的模型）
# TODO 需要先建立一个model文件夹
model_path = './model1/model.pth'
optimizer_path = './model1/optimizer.pth'


def train(epoch):
    network.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        # TODO 需要使用optimizer.zero_grad()手动将梯度设置为零，因为PyTorch在默认情况下会累积梯度
        optimizer.zero_grad()
        output = network(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset), 100. * batch_idx / len(train_loader), loss.item()))
            train_losses.append(loss.item())
            train_counter.append(
                (batch_idx * 64) + ((epoch - 1) * len(train_loader.dataset)))
    if epoch == (n_epochs - 1):
        # TODO 存储模型
        torch.save(network.state_dict(), model_path)
        torch.save(optimizer.state_dict(), optimizer_path)


def test():
    network.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            output = network(data)
            test_loss += F.nll_loss(output, target, size_average=False).item()
            pred = output.data.max(1, keepdim=True)[1]
            correct += pred.eq(target.data.view_as(pred)).sum()
    test_loss /= len(test_loader.dataset)
    test_losses.append(test_loss)
    print('\nTest set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))


for epoch in range(1, n_epochs + 1):
    train(epoch)
    test()

fig = plt.figure()
plt.plot(train_counter, train_losses, color='blue')
print(len(test_counter))
print(len(test_losses))
plt.scatter(test_counter, test_losses, color='red')
plt.legend(['Train Loss', 'Test Loss'], loc='upper right')
plt.xlabel('number of training examples seen')
plt.ylabel('negative log likelihood loss')
plt.show()