import os
os.environ['TF_CPP_MIN_LOG_LEVEL']='2'
import os.path
import sys
import cv2
import numpy as np
import json
import base64
import glob
import tensorflow as tf
from tensorflow.core.framework import graph_pb2

def _load_test_data(test_data_paths):
    features = { 'image/filename':   tf.FixedLenFeature((), tf.string, default_value=''),
                 'image/encoded':    tf.FixedLenFeature((), tf.string, default_value=''),
                 'image/object/bbox/xmin':  tf.FixedLenFeature((), tf.float32, default_value=0.),
                 'image/object/bbox/xmax':  tf.FixedLenFeature((), tf.float32, default_value=1.),
                 'image/object/bbox/ymin':  tf.FixedLenFeature((), tf.float32, default_value=0.),
                 'image/object/bbox/ymax':  tf.FixedLenFeature((), tf.float32, default_value=1.),
                 'image/object/class/label':      tf.FixedLenFeature((), tf.int64, default_value=0)}

    filename_queue = tf.train.string_input_producer(test_data_paths, num_epochs=1)
    reader = tf.TFRecordReader()
    _, example = reader.read(filename_queue)
    parsed_features = tf.parse_single_example(example, features=features)

    image =  tf.image.decode_jpeg(parsed_features['image/encoded'], channels = 3)

    label = tf.to_int32(parsed_features['image/object/class/label'])

    xmin = parsed_features['image/object/bbox/xmin']
    xmin = tf.to_int32(xmin * tf.to_float(tf.shape(image)[1]))
    xmax = parsed_features['image/object/bbox/xmax']
    xmax = tf.to_int32(xmax * tf.to_float(tf.shape(image)[1]))
    ymin = parsed_features['image/object/bbox/ymin']
    ymin = tf.to_int32(ymin * tf.to_float(tf.shape(image)[0]))
    ymax = parsed_features['image/object/bbox/ymax']
    ymax = tf.to_int32(ymax * tf.to_float(tf.shape(image)[0]))
    image = tf.image.crop_to_bounding_box(image, ymin, xmin, ymax - ymin - 1, xmax - xmin - 1)
    image = tf.reverse(image, [-1]) # RGB to BGR because we use BGR on inference enter
    return tf.expand_dims(image, axis = 0), label

def _load_graph(graph_path, image_ph):
    graph_def = graph_pb2.GraphDef()
    with open(graph_path, "rb") as f:
        graph_def.ParseFromString(f.read())

    for node in graph_def.node:
        #print(node.name, node.op, [input for input in node.input])
        if node.op == 'RefSwitch':
            node.op = 'Switch'
            for index in range(len(node.input)):
                if 'moving_' in node.input[index]:
                    node.input[index] = node.input[index] + '/read'
        elif node.op == 'AssignSub':
            node.op = 'Sub'
            if 'use_locking' in node.attr: del node.attr['use_locking']
        elif node.op == 'AssignAdd':
            node.op = 'Add'
            if 'use_locking' in node.attr: del node.attr['use_locking']

    tf.import_graph_def(graph_def, name='', input_map={'inference_input': image_ph})
    softmax_tensor = tf.get_default_graph().get_tensor_by_name('inference_softmax:0')
    return softmax_tensor

def calc_stat(graph_path, test_data_paths, base_thr = 0.96, digits_len = 5):
    NUM2WORDS = {1: 'one',
                 2: 'two',
                 3: 'three',
                 4: 'four',
                 5: 'five',
                 6: 'six',
                 7: 'seven'}

    image, label_gt = _load_test_data(test_data_paths)
    softmax_tensor = _load_graph(graph_path, image)

    conf_thr = base_thr ** digits_len
    errors = [0] * (digits_len + 1)
    all_conf_numbers = 0
    all_numbers = 0

    init_op = tf.group(tf.global_variables_initializer(),
                       tf.local_variables_initializer())

    with tf.Session()  as sess:
        sess.run(init_op)
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(coord=coord)
        try:
            idx  = 0
            while (True):
                number_gt, softmax_result = sess.run([label_gt, softmax_tensor])
                argmax = np.argmax(softmax_result[0], axis = -1)
                all_numbers += 1

                number = 0
                for i in range(argmax.shape[0]):
                    if argmax[i] >= 10:
                        continue # this can make problems for case then 10 in the middle
                                 # but this case has problem anyway
                    digit = argmax[i]
                    number = number * 10 + digit
                confidence = np.prod(np.max(softmax_result, axis = -1),  axis = -1)
                if (confidence < conf_thr):
                    continue

                number_test = int(number)
                number_gt = str(number_gt)
                number_test = str(number_test)
                errs_cnt = abs(len(number_gt) - len(number_test))
                for x, y in zip(reversed(number_gt), reversed(number_test)):
                    errs_cnt += x != y
                all_conf_numbers += 1
                errors[errs_cnt] += 1

        except tf.errors.OutOfRangeError:
            print("Test dataset ended")
        finally:
            coord.request_stop()
        coord.request_stop()
        coord.join(threads)

    if (0 < all_conf_numbers):
        result  = "Accuracy: {:.2f} %\n".format(errors[0] / all_conf_numbers * 100.)
        for i in range(1, digits_len + 1):
            result += "  with {} error: {:.2f} %\n".format(NUM2WORDS[i], errors[i] / all_conf_numbers * 100.)
        result += "Valid numbers {} from {} numbers with detection great than confidence\n".format(errors[0], all_conf_numbers)
        result += "{} numbers in dataset".format(all_numbers)
        return result
    else:
        return "NO RESULTS WITH CONFIDENCE MORE THAN {}".format(conf_thr)

