import time
import os
import math

import numpy as np

from libs import utils
from libs.img_dataset import ImgDataset
from nets.crnn import CRNN
from nets.cnn.paper_cnn import PaperCNN
import shutil


def calculate_accuracy(predicts, labels):
    """
    :param predicts: encoded predict result
    :param labels: ground true label
    :return: accuracy
    """
    assert len(predicts) == len(labels)

    correct_count = 0
    for i, p_label in enumerate(predicts):
        if p_label == labels[i]:
            correct_count += 1

    acc = correct_count / len(predicts)
    return acc, correct_count


def calculate_edit_distance_mean(edit_distences):
    """
    排除了 edit_distance == 0 的值计算编辑距离的均值
    :param edit_distences:
    :return:
    """
    data = np.array(edit_distences)
    data = data[data != 0]
    if len(data) == 0:
        return 0
    return np.mean(data)


def validation(sess, feeds, fetches, dataset, converter, result_dir, name,
               step=None, print_batch_info=False, copy_failed=False):
    """
    Save file name: {acc}_{step}.txt
    :param sess: tensorflow session
    :param model: crnn network
    :param result_dir:
    :param name: val, test, infer. used to create sub dir in result_dir
    :return:
    """
    sess.run(dataset.init_op)

    img_paths = []
    predicts = []
    trimed_predicts = []
    labels = []
    trimed_labels = []
    edit_distances = []
    total_batch_time = 0

    for batch in range(dataset.num_batches):
        img_batch, widths, label_batch, batch_labels, batch_img_paths = dataset.get_next_batch(sess)
        if len(batch_labels) == 0:
            continue

        batch_start_time = time.time()

        feed = {feeds['inputs']: img_batch,
                feeds['labels']: label_batch,
                feeds['sequence_length']: PaperCNN.get_sequence_lengths(widths),
                feeds['is_training']: False}

        try:
            batch_predicts, edit_distance, batch_edit_distances = sess.run(fetches, feed)
        except Exception:
            print(batch_labels)
            continue
        batch_predicts = [converter.decode(p, CRNN.CTC_INVALID_INDEX) for p in batch_predicts]

        trimed_batch_predicts = [utils.remove_all_symbols(txt) for txt in batch_predicts]
        trimed_batch_labels = [utils.remove_all_symbols(txt) for txt in batch_labels]

        img_paths.extend(batch_img_paths)
        predicts.extend(batch_predicts)
        labels.extend(batch_labels)
        trimed_predicts.extend(trimed_batch_predicts)
        trimed_labels.extend(trimed_batch_labels)
        edit_distances.extend(batch_edit_distances)

        acc, correct_count = calculate_accuracy(batch_predicts, batch_labels)
        trimed_acc, trimed_correct_count = calculate_accuracy(trimed_batch_predicts, trimed_batch_labels)

        batch_time = time.time() - batch_start_time
        total_batch_time += batch_time
        if print_batch_info:
            print("{:.03f}s [{}/{}] acc: {:.03f}({}/{}), edit_distance: {:.03f}, trim_acc {:.03f}({}/{})"
                  .format(batch_time, batch, dataset.num_batches,
                          acc, correct_count, dataset.batch_size,
                          edit_distance,
                          trimed_acc, trimed_correct_count, dataset.batch_size))

    acc, correct_count = calculate_accuracy(predicts, labels)
    trimed_acc, trimed_correct_count = calculate_accuracy(trimed_predicts, trimed_labels)
    edit_distance_mean = calculate_edit_distance_mean(edit_distances)
    total_edit_distance = sum(edit_distances)

    acc_str = "Accuracy: {:.03f} ({}/{}), Trimed Accuracy: {:.03f} ({}/{})" \
              "Total edit distance: {:.03f}, " \
              "Average edit distance: {:.03f}, Average batch time: {:.03f}" \
        .format(acc, correct_count, dataset.size,
                trimed_acc, trimed_correct_count, dataset.size,
                total_edit_distance, edit_distance_mean, total_batch_time / dataset.num_batches)

    print(acc_str)

    save_dir = os.path.join(result_dir, name)
    utils.check_dir_exist(save_dir)

    result_file_path = save_txt_result(save_dir, acc, step, labels, predicts, 'acc',
                                       edit_distances, acc_str)

    save_txt_result(save_dir, acc, step, labels, predicts, 'acc', edit_distances,
                    acc_str, only_failed=True)

    save_txt_result(save_dir, trimed_acc, step, trimed_labels, trimed_predicts, 'tacc',
                    edit_distances)

    save_txt_result(save_dir, trimed_acc, step, trimed_labels, trimed_predicts, 'tacc',
                    edit_distances, only_failed=True)

    save_txt_4_analyze(save_dir, labels, predicts, 'acc', step)
    save_txt_4_analyze(save_dir, trimed_labels, trimed_predicts, 'tacc', step)

    # Copy image not all match to a dir
    # TODO: we will only save failed imgs for acc
    if copy_failed:
        failed_infer_img_dir = result_file_path[:-4] + "_failed"
        if os.path.exists(failed_infer_img_dir) and os.path.isdir(failed_infer_img_dir):
            shutil.rmtree(failed_infer_img_dir)

        utils.check_dir_exist(failed_infer_img_dir)

        failed_image_indices = []
        for i, val in enumerate(edit_distances):
            if val != 0:
                failed_image_indices.append(i)

        for i in failed_image_indices:
            img_path = img_paths[i]
            img_name = img_path.split("/")[-1]
            dst_path = os.path.join(failed_infer_img_dir, img_name)
            shutil.copyfile(img_path, dst_path)

        failed_infer_result_file_path = os.path.join(failed_infer_img_dir, "result.txt")
        with open(failed_infer_result_file_path, 'w', encoding='utf-8') as f:
            for i in failed_image_indices:
                p_label = predicts[i]
                t_label = labels[i]
                f.write("{}\n".format(img_paths[i]))
                f.write("input:   {:17s} length: {}\n".format(t_label, len(t_label)))
                f.write("predict: {:17s} length: {}\n".format(p_label, len(p_label)))
                f.write("edit distance:  {}\n".format(edit_distances[i]))
                f.write('-' * 30 + '\n')

    return acc, trimed_acc, edit_distance_mean, total_edit_distance, correct_count, trimed_correct_count


def save_txt_4_analyze(save_dir, labels, predicts, acc_type, step):
    """
    把测试集的真值和预测结果放在保存在同一个 txt 文件中，方便统计
    """
    txt_path = os.path.join(save_dir, '%d_%s_gt_and_pred.txt' % (step, acc_type))

    with open(txt_path, 'w', encoding='utf-8') as f:
        for i, p_label in enumerate(predicts):
            t_label = labels[i]
            f.write("{}__$__{}\n".format(t_label, p_label))


def save_txt_result(save_dir, acc, step, labels, predicts, acc_type,
                    edit_distances=None, acc_str=None, only_failed=False):
    """
    :param acc_type:  'acc' or 'tacc'
    :return:
    """
    failed_suffix = ''
    if only_failed:
        failed_suffix = 'failed'

    if step is not None:
        txt_path = os.path.join(save_dir, '%d_%s_%.3f_%s.txt' % (step, acc_type, acc, failed_suffix))
    else:
        txt_path = os.path.join(save_dir, '%s_%.3f_%s.txt' % (acc_type, acc, failed_suffix))

    print("Write result to %s" % txt_path)
    with open(txt_path, 'w', encoding='utf-8') as f:
        for i, p_label in enumerate(predicts):
            t_label = labels[i]
            all_match = (t_label == p_label)

            if only_failed and all_match:
                continue

            # f.write("{}\n".format(img_paths[i]))
            f.write("input:   {:17s} length: {}\n".format(t_label, len(t_label)))
            f.write("predict: {:17s} length: {}\n".format(p_label, len(p_label)))
            f.write("all match:  {}\n".format(1 if all_match else 0))

            if edit_distances:
                f.write("edit distance:  {}\n".format(edit_distances[i]))

            f.write('-' * 30 + '\n')

        if acc_str:
            f.write(acc_str + "\n")

    return txt_path

[((37, 11, 37, 35), 'numpy.array', 'np.array', ({(37, 20, 37, 34): 'edit_distences'}, {}), '(edit_distences)', True, 'import numpy as np\n'), ((41, 11, 41, 24), 'numpy.mean', 'np.mean', ({(41, 19, 41, 23): 'data'}, {}), '(data)', True, 'import numpy as np\n'), ((119, 15, 119, 45), 'os.path.join', 'os.path.join', ({(119, 28, 119, 38): 'result_dir', (119, 40, 119, 44): 'name'}, {}), '(result_dir, name)', False, 'import os\n'), ((120, 4, 120, 35), 'libs.utils.check_dir_exist', 'utils.check_dir_exist', ({(120, 26, 120, 34): 'save_dir'}, {}), '(save_dir)', False, 'from libs import utils\n'), ((175, 15, 175, 81), 'os.path.join', 'os.path.join', ({(175, 28, 175, 36): 'save_dir', (175, 38, 175, 80): "'%d_%s_gt_and_pred.txt' % (step, acc_type)"}, {}), "(save_dir, '%d_%s_gt_and_pred.txt' % (step, acc_type))", False, 'import os\n'), ((69, 27, 69, 38), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((144, 8, 144, 51), 'libs.utils.check_dir_exist', 'utils.check_dir_exist', ({(144, 30, 144, 50): 'failed_infer_img_dir'}, {}), '(failed_infer_img_dir)', False, 'from libs import utils\n'), ((157, 40, 157, 88), 'os.path.join', 'os.path.join', ({(157, 53, 157, 73): 'failed_infer_img_dir', (157, 75, 157, 87): '"""result.txt"""'}, {}), "(failed_infer_img_dir, 'result.txt')", False, 'import os\n'), ((194, 19, 194, 101), 'os.path.join', 'os.path.join', ({(194, 32, 194, 40): 'save_dir', (194, 42, 194, 100): "'%d_%s_%.3f_%s.txt' % (step, acc_type, acc, failed_suffix)"}, {}), "(save_dir, '%d_%s_%.3f_%s.txt' % (step, acc_type, acc,\n    failed_suffix))", False, 'import os\n'), ((196, 19, 196, 92), 'os.path.join', 'os.path.join', ({(196, 32, 196, 40): 'save_dir', (196, 42, 196, 91): "'%s_%.3f_%s.txt' % (acc_type, acc, failed_suffix)"}, {}), "(save_dir, '%s_%.3f_%s.txt' % (acc_type, acc, failed_suffix))", False, 'import os\n'), ((73, 42, 73, 79), 'nets.cnn.paper_cnn.PaperCNN.get_sequence_lengths', 'PaperCNN.get_sequence_lengths', ({(73, 72, 73, 78): 'widths'}, {}), '(widths)', False, 'from nets.cnn.paper_cnn import PaperCNN\n'), ((83, 33, 83, 62), 'libs.utils.remove_all_symbols', 'utils.remove_all_symbols', ({(83, 58, 83, 61): 'txt'}, {}), '(txt)', False, 'from libs import utils\n'), ((84, 31, 84, 60), 'libs.utils.remove_all_symbols', 'utils.remove_all_symbols', ({(84, 56, 84, 59): 'txt'}, {}), '(txt)', False, 'from libs import utils\n'), ((96, 21, 96, 32), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((141, 11, 141, 47), 'os.path.exists', 'os.path.exists', ({(141, 26, 141, 46): 'failed_infer_img_dir'}, {}), '(failed_infer_img_dir)', False, 'import os\n'), ((141, 52, 141, 87), 'os.path.isdir', 'os.path.isdir', ({(141, 66, 141, 86): 'failed_infer_img_dir'}, {}), '(failed_infer_img_dir)', False, 'import os\n'), ((142, 12, 142, 47), 'shutil.rmtree', 'shutil.rmtree', ({(142, 26, 142, 46): 'failed_infer_img_dir'}, {}), '(failed_infer_img_dir)', False, 'import shutil\n'), ((154, 23, 154, 67), 'os.path.join', 'os.path.join', ({(154, 36, 154, 56): 'failed_infer_img_dir', (154, 58, 154, 66): 'img_name'}, {}), '(failed_infer_img_dir, img_name)', False, 'import os\n'), ((155, 12, 155, 47), 'shutil.copyfile', 'shutil.copyfile', ({(155, 28, 155, 36): 'img_path', (155, 38, 155, 46): 'dst_path'}, {}), '(img_path, dst_path)', False, 'import shutil\n')]