# -*- coding: UTF-8 -*-
import numpy as np
import torch
from torch.autograd import Variable
import captcha_setting
import my_dataset
from captcha_cnn_model import CNN
import one_hot_encoding
import matplotlib.pyplot as plt


def main():
    cnn = CNN()
    cnn.eval()
    cnn.load_state_dict(torch.load('model.pkl'))
    print("load cnn net.")

    test_dataloader = my_dataset.get_test_data_loader()
    #print(len(list(enumerate(test_dataloader))))
    correct = 0
    total = 0
    for i, (images, labels) in enumerate(test_dataloader):
        image = images
        # convert_image = image[0, 0, :, :]
        # plt.imshow(convert_image)
        vimage = Variable(image)
        predict_label = cnn(vimage)
        c0 = captcha_setting.ALL_CHAR_SET[np.argmax(predict_label[0, 0:captcha_setting.ALL_CHAR_SET_LEN].data.numpy())]
        c1 = captcha_setting.ALL_CHAR_SET[np.argmax(predict_label[0, captcha_setting.ALL_CHAR_SET_LEN:2 * captcha_setting.ALL_CHAR_SET_LEN].data.numpy())]
        c2 = captcha_setting.ALL_CHAR_SET[np.argmax(predict_label[0, 2 * captcha_setting.ALL_CHAR_SET_LEN:3 * captcha_setting.ALL_CHAR_SET_LEN].data.numpy())]
        c3 = captcha_setting.ALL_CHAR_SET[np.argmax(predict_label[0, 3 * captcha_setting.ALL_CHAR_SET_LEN:4 * captcha_setting.ALL_CHAR_SET_LEN].data.numpy())]
        predict_label = '%s%s%s%s' % (c0, c1, c2, c3)
        true_label = one_hot_encoding.decode(labels.numpy()[0])
        # plt.title("predict: "+predict_label +  " true:" + true_label,fontsize=20,color='red')
        # plt.show()
        print(true_label," // ",predict_label)
        total += labels.size(0)
        if(predict_label == true_label):
            correct += 1
        if(total%1000==0):
            print('Test Accuracy of the model on the %d test images: %f %%' % (total, 100 * correct / total))
        # if (i > 5):
        #     break
    print('correct_num=%d  ---  total_num = %d ' % (correct,total))
    print('Test Accuracy of the model on the %d test images: %f %%' % (total, 100 * correct / total))

if __name__ == '__main__':
    main()