import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.utils.data as Data
import torchvision
import matplotlib.pyplot as plt

#Hyper prameters
EPOCH = 2
BATCH_SIZE = 50
LR = 0.001
DOWNLOAD_MNIST = True
train_data = torchvision.datasets.MNIST(
    root ='./mnist',
    train = True,
    download = DOWNLOAD_MNIST
)

print(train_data.data.size())
print(train_data.targets.size())
print(train_data.data[0])

# 畫一個圖片顯示出來
plt.imshow(train_data.data[0].numpy(),cmap='gray')
plt.title('%i'%train_data.targets[0])
plt.show()

#訓練和測試數據準備
train_loader=Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)
test_data=torchvision.datasets.MNIST(
    root='./mnist',
    train=False,
)
#這里只取前3千個數據吧，差不多已經夠用了，然后將其歸一化。
with torch.no_grad():
    test_x=Variable(torch.unsqueeze(test_data.data, dim=1)).type(torch.FloatTensor)[:3000]/255   
    test_y=test_data.targets[:3000]

'''開始建立CNN網絡'''
class CNN(nn.Module):
    def __init__(self):
        super(CNN,self).__init__()
        '''
        一般來說，卷積網絡包括以下內容：
        1.卷積層
        2.神經網絡
        3.池化層
        '''
        self.conv1=nn.Sequential(
            nn.Conv2d(              #--> (1,28,28)
                in_channels=1,      #傳入的圖片是幾層的，灰色為1層，RGB為三層
                out_channels=16,    #輸出的圖片是幾層
                kernel_size=5,      #代表掃描的區域點為5*5
                stride=1,           #就是每隔多少步跳一下
                padding=2,          #邊框補全，其計算公式=（kernel_size-1）/2=(5-1)/2=2
            ),    # 2d代表二維卷積           --> (16,28,28)
            nn.ReLU(),              #非線性激活層
            nn.MaxPool2d(kernel_size=2),    #設定這里的掃描區域為2*2，且取出該2*2中的最大值          --> (16,14,14)
        )
        self.conv2=nn.Sequential(
            nn.Conv2d(              #       --> (16,14,14)
                in_channels=16,     #這里的輸入是上層的輸出為16層
                out_channels=32,    #在這里我們需要將其輸出為32層
                kernel_size=5,      #代表掃描的區域點為5*5
                stride=1,           #就是每隔多少步跳一下
                padding=2,          #邊框補全，其計算公式=（kernel_size-1）/2=(5-1)/2=
            ),                      #   --> (32,14,14)
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),    #設定這里的掃描區域為2*2，且取出該2*2中的最大值     --> (32,7,7)，這里是三維數據
        )
        self.out=nn.Linear(32*7*7,10)       #注意一下這里的數據是二維的數據
    def forward(self,x):
        x=self.conv1(x)
        x=self.conv2(x)     #（batch,32,7,7）
        #然后接下來進行一下擴展展平的操作，將三維數據轉為二維的數據
        x=x.view(x.size(0),-1)    #(batch ,32 * 7 * 7)
        output=self.out(x)
        return output

cnn=CNN()
print(cnn)

# 添加優化方法
optimizer=torch.optim.Adam(cnn.parameters(),lr=LR)
# 指定損失函數使用交叉信息熵
loss_fn=nn.CrossEntropyLoss()
'''
開始訓練我們的模型哦
'''
step=0
for epoch in range(EPOCH):
    #加載訓練數據
    for step,data in enumerate(train_loader):
        x,y=data
        #分別得到訓練數據的x和y的取值
        b_x=Variable(x)
        b_y=Variable(y)
        output=cnn(b_x)         #調用模型預測
        loss=loss_fn(output,b_y)#計算損失值
        optimizer.zero_grad()   #每一次循環之前，將梯度清零
        loss.backward()         #反向傳播
        optimizer.step()        #梯度下降
        #每執行50次，輸出一下當前epoch、loss、accuracy
        if (step%50==0):
            #計算一下模型預測正確率
            test_output=cnn(test_x)
            y_pred=torch.max(test_output,1)[1].data.squeeze()
            accuracy=sum(y_pred==test_y).item()/test_y.size(0)
            print('now epoch :  ', epoch, '   |  loss : %.4f ' % loss.item(), '     |   accuracy :   ' , accuracy)
'''
打印十個測試集的結果
'''
test_output=cnn(test_x[:10])
y_pred=torch.max(test_output,1)[1].data.squeeze()       #選取最大可能的數值所在的位置
print(y_pred.tolist(),'predecton Result')
print(test_y[:10].tolist(),'Real Result')