import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
import torchvision
from torch.autograd import Variable
from torch.utils.data import DataLoader
import cv2

# 下載訓練集
train_dataset = datasets.MNIST(root='E:\mnist',
                               train=True,
                               transform=transforms.ToTensor(),
                               download=True)
# 下載測試集
test_dataset = datasets.MNIST(root='E:\mnist',
                              train=False,
                              transform=transforms.ToTensor(),
                              download=True)

# dataset 參數用于指定我們載入的數據集名稱
# batch_size參數設置了每個包中的圖片數據個數
# 在裝載的過程會將數據隨機打亂順序并進打包
batch_size = 64
# 建立一個數據迭代器
# 裝載訓練集
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           shuffle=True)
# 裝載測試集
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size,
                                          shuffle=True)


# 卷積層使用 torch.nn.Conv2d
# 激活層使用 torch.nn.ReLU
# 池化層使用 torch.nn.MaxPool2d
# 全連接層使用 torch.nn.Linear
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Sequential(nn.Conv2d(1, 6, 3, 1, 2),
                                   nn.ReLU(), nn.MaxPool2d(2, 2))

        self.conv2 = nn.Sequential(nn.Conv2d(6, 16, 5), nn.ReLU(),
                                   nn.MaxPool2d(2, 2))

        self.fc1 = nn.Sequential(nn.Linear(16 * 5 * 5, 120),
                                 nn.BatchNorm1d(120), nn.ReLU())

        self.fc2 = nn.Sequential(
            nn.Linear(120, 84),
            nn.BatchNorm1d(84),
            nn.ReLU(),
            nn.Linear(84, 10))
        # 最后的結果一定要變為 10，因為數字的選項是 0 ~ 9

    def forward(self, x):
        x = self.conv1(x)
        # print("1:", x.shape)
        # 1: torch.Size([64, 6, 30, 30])
        # max pooling
        # 1: torch.Size([64, 6, 15, 15])
        x = self.conv2(x)
        # print("2:", x.shape)
        # 2: torch.Size([64, 16, 5, 5])
        # 對參數實現扁平化
        x = x.view(x.size()[0], -1)
        x = self.fc1(x)
        x = self.fc2(x)
        return x


def test_image_data(images, labels):
    # 初始輸出為一段數字圖像序列
    # 將一段圖像序列整合到一張圖片上 (make_grid會默認將圖片變成三通道，默認值為0)
    # images: torch.Size([64, 1, 28, 28])
    img = torchvision.utils.make_grid(images)
    # img: torch.Size([3, 242, 242])
    # 將通道維度置在第三個維度
    img = img.numpy().transpose(1, 2, 0)
    # img: torch.Size([242, 242, 3])
    # 減小圖像對比度
    std = [0.5, 0.5, 0.5]
    mean = [0.5, 0.5, 0.5]
    img = img * std + mean
    # print(labels)
    cv2.imshow('win2', img)
    key_pressed = cv2.waitKey(0)


# 初始化設備信息
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 學習速率
LR = 0.001
# 初始化網絡
net = LeNet().to(device)
# 損失函數使用交叉熵
criterion = nn.CrossEntropyLoss()
# 優化函數使用 Adam 自適應優化算法
optimizer = optim.Adam(net.parameters(), lr=LR, )
epoch = 1
if __name__ == '__main__':
    for epoch in range(epoch):
        print("GPU:", torch.cuda.is_available())
        sum_loss = 0.0
        for i, data in enumerate(train_loader):
            inputs, labels = data
            # print(inputs.shape)
            # torch.Size([64, 1, 28, 28])
            # 將內存中的數據復制到gpu顯存中去
            inputs, labels = Variable(inputs).cuda(), Variable(labels).cuda()
            # 將梯度歸零
            optimizer.zero_grad()
            # 將數據傳入網絡進行前向運算
            outputs = net(inputs)
            # 得到損失函數
            loss = criterion(outputs, labels)
            # 反向傳播
            loss.backward()
            # 通過梯度做一步參數更新
            optimizer.step()
            # print(loss)
            sum_loss += loss.item()
            if i % 100 == 99:
                print('[%d,%d] loss:%.03f' % (epoch + 1, i + 1, sum_loss / 100))
                sum_loss = 0.0
                # 將模型變換為測試模式
        net.eval()
        correct = 0
        total = 0
        for data_test in test_loader:
            _images, _labels = data_test
            # 將內存中的數據復制到gpu顯存中去
            images, labels = Variable(_images).cuda(), Variable(_labels).cuda()
            # 圖像預測結果
            output_test = net(images)
            # torch.Size([64, 10])
            # 從每行中找到最大預測索引
            _, predicted = torch.max(output_test, 1)
            # 圖像可視化
            # print("predicted:", predicted)
            # test_image_data(_images, _labels)
            # 預測數據的數量
            total += labels.size(0)
            # 預測正確的數量
            correct += (predicted == labels).sum()
        print("correct1: ", correct)
        print("Test acc: {0}".format(correct.item() / total))