import os
import argparse
import json
from .methods import calculate
import json
import sys
# from hikframe.dl.cv.utils.online_infer import infer_ret_parser
# cadicate = {
#     'img_classification': ['precision', 'recall', 'F1-Score', 'acc', 'mAP'],
#     'img_detection': ['precision', 'recall', 'F1-Score', 'mAP']
# }
cadicate = {
    'img_classification': ["Accuracy", "Macro-Precision", "Macro-Recall", "Macro-F1-Score"],
    'img_detection': ["mAP", "Macro-Precision", "Macro-Recall", "Macro-F1-Score"]
}


def infer_ret_parser(det_ret):
    img_names = list(det_ret.keys())
    # 输出拼接
    rets = []
    for img_name in img_names:
        ret = ''
        ret += img_name
        for det in det_ret[img_name]:
            if len(det) == 6:
                label, score, x_min, y_min, x_max, y_max = det
                ret += ' '
                pred = '{}:{:.3f} {} {} {} {}'.format(label, score, x_min,
                                                      y_min, x_max, y_max)
                ret += pred
            elif len(det) == 2:
                label, score = det
                ret += ' '
                pred = '{}:{:.3f}'.format(label, score)
                ret += pred
            else:
                raise NotImplementedError
        rets.append(ret)
    return rets


def ret_generater(rets, txt_path):
    if type(rets) == dict:
        rets = infer_ret_parser(rets)
    with open(txt_path, 'a') as f:
        for line in rets:
            f.writelines(line + '\n')
            print(line)


def eval(ret, anno_path, save_dir, type, methods, scoreThred, IOUThred=None, process_num=20):
    predict_txt_name = 'predict.txt'
    predict_path = os.path.join(save_dir, predict_txt_name)

    ret_generater(ret, txt_path=predict_path)

    if not os.path.isfile(predict_path):
        print('预测文件不存在，请确认位置正确')
        return
    if not os.path.exists(anno_path):
        print('标签目录不存在，请确认数据集位置正确')
        return

    assert type in cadicate, 'wrong type'
    calculate(type, methods, scoreThred, IOUThred,
              predict_path, anno_path, save_dir, process_num)
    return


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='Simple training script for training a RetinaNet network.')

    parser.add_argument('--type', type=str, help="task type")
    parser.add_argument('--method', type=str, help="eval methods")
    parser.add_argument('--predict_path', type=str, help="path to predict txt")
    parser.add_argument('--anno_path', type=str, help="path to anno_path")
    parser.add_argument('--anno_path', type=str, help="path to anno_path")
    parser.add_argument('--scoreThred', type=str, help="scoreThred")
    parser.add_argument('--IOUThred', type=str, help="IOUThred", default='')

    parser = parser.parse_args()

    eval(parser.predict_path, parser.anno_path, parser.type,
         parser.methods, parser.scoreThred, parser.IOUThred)
