import matplotlib.pyplot as plt

import argparse

import csv

import json

import os

from collections import namedtuple

from PIL import Image

import editdistance

import mxnet as mx

import numpy as np

from callbacks.save_bboxes import BBOXPlotter

from metrics.ctc_metrics import strip_prediction

from networks.text_rec import SVHNMultiLineCTCNetwork

from operations.disable_shearing import *

from utils.datatypes import Size

Batch = namedtuple('Batch', ['data'])

#后缀都不能加的，程序自己添加，似乎同时加载两个文件

sym,arg_params,aux_params = mx.model.load_checkpoint('./testxt/model/model',2)

#这里面应该是训练的参数

#print(arg_params)

net, loc, transformed_output, size_params = SVHNMultiLineCTCNetwork.get_network((1,1,64,200),Size(50,50),46,2,23)

output = mx.sym.Group([loc, transformed_output, net])

#靠 在这里预定义的话，TMD，soft 层怎么办？

mod = mx.mod.Module(output,context=mx.cpu(),data_names=['data',

        'softmax_label',

	'l0_forward_init_h_state',

	'l0_forward_init_c_state_cell',

	'l1_forward_init_h_state',

	'l1_forward_init_c_state_cell'

],label_names=[])

mod.bind(for_training=False,grad_req='null',data_shapes=[

	    ('data',(1,1,64,200)),

            ('softmax_label', (1,23)),

	    ('l0_forward_init_h_state', (1, 1, 256)),

            ('l0_forward_init_c_state_cell', (1, 1, 256)),

            ('l1_forward_init_h_state', (1, 1, 256)),

            ('l1_forward_init_c_state_cell', (1, 1, 256))

	])

arg_params['l0_forward_init_h_state'] = mx.nd.zeros((1, 1, 256))

arg_params['l0_forward_init_c_state_cell'] = mx.nd.zeros((1, 1, 256))

arg_params['l1_forward_init_h_state'] = mx.nd.zeros((1, 1, 256))

arg_params['l1_forward_init_c_state_cell'] = mx.nd.zeros((1, 1, 256))

mod.set_params(arg_params, aux_params)

#看看怎么加载label

#一个映射文件，类似caffe中的label,在下面循环中用到

with open('/home/lbk/python-env/stn-ocr/mxnet/testxt/ctc_char_map.json') as char_map_file:

    char_map = json.load(char_map_file)

reverse_char_map = {v: k for k, v in char_map.items()}

print(len(reverse_char_map))

with open('/home/lbk/python-env/stn-ocr/mxnet/testxt/icdar2013_eval/one_gt.txt') as eval_gt:

    reader = csv.reader(eval_gt,delimiter=';')

    for idx,line in enumerate(reader):

        file_name = line[0]

        label = line[1].strip()

        gt_word = label.lower()

        print(gt_word)

        #这一步又是干什么的

        #dict.get(key,default)查找，不存在返回default

        label = [reverse_char_map.get(ord(char.lower()),reverse_char_map[9250]) for char in gt_word]

        label+=[reverse_char_map[9250]]*(23-len(label))

        #print(label)

        the_image = Image.open(file_name)

        the_image = the_image.convert('L')

        the_image = the_image.resize((200,64), Image.ANTIALIAS)

        image = np.asarray(the_image, dtype=np.float32)[np.newaxis, np.newaxis, ...]

        image/=255

        temp = mx.nd.zeros((1,1,256))

        label = mx.nd.array([label])

        image = mx.nd.array(image)

        print(type(temp),type(label))

        input_batch = Batch(data=[image,label,temp,temp,temp,temp])

        mod.forward(input_batch,is_train=False)

        print(len(mod.get_outputs()))

        print('0000',mod.get_outputs()[2])

        predictions = mod.get_outputs()[2].asnumpy()

        predicted_classes = np.argmax(predictions,axis=1)

        print(len(predicted_classes))

        print(predicted_classes)

        predicted_classes = strip_prediction(predicted_classes, int(reverse_char_map[9250]))

        predicted_word = ''.join([chr(char_map[str(p)]) for p in predicted_classes]).replace(' ', '')

        print(predicted_word)

        distance = editdistance.eval(gt_word, predicted_word)

        print("{} - {}\t\t{}: {}".format(idx, gt_word, predicted_word, distance))

        results = [prediction == label for prediction, label in zip(predicted_word, gt_word)]

        print(results)