import pdb
import os
import sys
import json
import shutil
import xml.etree.ElementTree as ET
import subprocess
import json
from copy import deepcopy
from multiprocessing import Process, Manager
# print(__file__)
# sys.path.append(os.path.dirname(__file__))
# print(os.path.dirname(__file__))
CUR_DIR = os.path.dirname(__file__)

g_process_num = 20

if not (os.getenv('hikmars_version') and os.getenv('hikmars_version') == '3'):
    try:
        from hikframe.dl.base.evaluate import Evaluation
        EL = Evaluation()
    except:
        pass
    print('hikmars_version: 4')
    is_hikmars3 = False
else:
    is_hikmars3 = True
    print('hikmars_version: 3')


def det_anno_parser(anno_path, gt_dir, results_dir, format='voc'):
    results_dict_template = {
        "img_file": "",
        "img_size": {
            "height": "",
            "width": "",
            "depth": ""
        },
        "gt": {
            "object": []
        },
        "predicted": {
            "detected_object": []
        },
        "info": ""
    }
    if format == 'coco':
        anno = os.path.join(anno_path, 'coco.json')
        infos = {}
        with open(anno, 'r') as f:
            infos = json.load(f)
        if not infos:
            print('ERROR: coco json parser wrong, please check it.')
            return
        categories = infos['categories']
        map_id2cate = {}
        for cat in categories:
            map_id2cate[cat["id"]] = cat["name"]

        map_id2img = {}
        images = infos['images']
        for img in images:
            file_name = img['file_name']
            height = img['height']
            width = img['width']
            id = img['id']
            map_id2img[id] = [file_name, height, width, []]
        annos = infos['annotations']
        for anno in annos:
            bbox = anno['bbox']
            category_id = anno['category_id']
            image_id = anno['image_id']
            img = map_id2img[image_id]
            cate = map_id2cate[category_id]
            map_id2img[image_id][3].append((cate, bbox))

        # multiprocess
        def _process(data, gt_dir, results_dir, results_dict_template):
            for info in data:
                results_dict = deepcopy(results_dict_template)
                print(info)
                file_name, height, width, boxes = info
                results_dict["img_size"]["height"] = height
                results_dict["img_size"]["width"] = width
                results_dict["img_size"]["depth"] = 3

                name = os.path.splitext(os.path.basename(file_name))[0]
                anno_txt = os.path.join(gt_dir, '{}.txt'.format(name))

                with open(anno_txt, 'a') as a_f:
                    index = -1
                    for box in boxes:
                        index += 1
                        t_obj = {
                            "label": "",
                            "bbox": {
                                "xmin": 0,
                                "ymin": 0,
                                "xmax": 0,
                                "ymax": 0
                            },
                            "index": 0
                        }
                        t_label = box[0]
                        xmin, ymin, b_width, b_height = box[1]
                        xmax, ymax = xmin+b_width, ymin+b_height
                        info = '{} {} {} {} {} {}'.format(t_label, xmin, ymin, xmax,
                                                          ymax, index) + '\n'
                        a_f.writelines(info)
                        t_obj["label"] = t_label
                        t_obj["bbox"]["xmin"] = xmin
                        t_obj["bbox"]["ymin"] = ymin
                        t_obj["bbox"]["xmax"] = xmax
                        t_obj["bbox"]["ymax"] = ymax
                        t_obj["index"] = index
                        results_dict["gt"]["object"].append(t_obj)
                info_json = os.path.join(results_dir, '{}.json'.format(name))
                with open(info_json, 'w') as f:
                    json.dump(results_dict, f)

        data = list(map_id2img.values())
        # multiprocess
        length = len(data)
        step = int(length / g_process_num) + 1

        process_list = []
        for i in range(0, length, step):  # 开启多个子进程执行fun1函数
            p = Process(target=_process, args=(
                data[i:i+step], gt_dir, results_dir, results_dict_template,))  # 实例化进程对象
            p.start()
            process_list.append(p)
        for p in process_list:
            p.join()

    elif format == 'voc':
        def _process(data, anno_path, gt_dir, results_dir, results_dict_template):
            for xml in data:
                results_dict = deepcopy(results_dict_template)
                if not xml.endswith('.xml'):
                    continue

                xml_path = os.path.join(anno_path, xml)
                tree = ET.parse(xml_path)
                root = tree.getroot()
                # name = root.find('filename').text.split('.')[0]
                # name = '.'.join(root.find('filename').text.split('.')[:-1])
                name = xml.split('.xml')[0]
                size = root.find('size')
                if size:
                    results_dict["img_size"]["height"] = size.find(
                        'height').text
                    results_dict["img_size"]["width"] = size.find('width').text
                    results_dict["img_size"]["depth"] = size.find('depth').text
                else:
                    IMG_HEIGHT = os.getenv('img_height')
                    IMG_WIDTH = os.getenv('img_width')
                    IMG_DEPTH = os.getenv('img_depth')
                    if IMG_HEIGHT and IMG_WIDTH and IMG_DEPTH:
                        results_dict["img_size"]["height"] = IMG_HEIGHT
                        results_dict["img_size"]["width"] = IMG_WIDTH
                        results_dict["img_size"]["depth"] = IMG_DEPTH
                    else:
                        print('no size find in xml of [{}], ignore this sample.'.format(
                            xml_path))
                        continue

                anno_txt = os.path.join(gt_dir, '{}.txt'.format(name))

                with open(anno_txt, 'a') as a_f:
                    index = -1
                    for obj in root.iter('object'):
                        index += 1
                        t_obj = {
                            "label": "",
                            "bbox": {
                                "xmin": 0,
                                "ymin": 0,
                                "xmax": 0,
                                "ymax": 0
                            },
                            "index": 0
                        }
                        t_label = obj.find('name').text
                        bndbox = obj.find('bndbox')
                        xmin = bndbox.find('xmin').text
                        ymin = bndbox.find('ymin').text
                        xmax = bndbox.find('xmax').text
                        ymax = bndbox.find('ymax').text
                        # info = '{} {} {} {} {}'.format(t_label, xmin, ymin, xmax,
                        #                                ymax) + '\n'
                        info = '{} {} {} {} {} {}'.format(t_label, xmin, ymin, xmax,
                                                          ymax, index) + '\n'
                        a_f.writelines(info)
                        t_obj["label"] = t_label
                        t_obj["bbox"]["xmin"] = xmin
                        t_obj["bbox"]["ymin"] = ymin
                        t_obj["bbox"]["xmax"] = xmax
                        t_obj["bbox"]["ymax"] = ymax
                        t_obj["index"] = index
                        results_dict["gt"]["object"].append(t_obj)
                info_json = os.path.join(results_dir, '{}.json'.format(name))
                with open(info_json, 'w') as f:
                    json.dump(results_dict, f)
        data = os.listdir(anno_path)
        # multiprocess
        length = len(data)
        step = int(length / g_process_num) + 1

        process_list = []
        for i in range(0, length, step):  # 开启多个子进程执行fun1函数
            p = Process(target=_process, args=(
                data[i:i+step], anno_path, gt_dir, results_dir, results_dict_template,))  # 实例化进程对象
            p.start()
            process_list.append(p)
        for p in process_list:
            p.join()


def clf_anno_parser(anno_path, gt_dir, results_dir, format='voc'):
    results_dict_template = {
        "img_file": "",
        "img_size": {
            "height": "",
            "width": "",
            "depth": ""
        },
        "gt": {
            "label": ""
        },
        "predicted": {
            "predicted_label": "",
            "scores": {}
        },
        "info": ""
    }
    if format == 'coco':
        raise NotImplementedError
    elif format == 'voc':
        def _process(data, anno_path, gt_dir, results_dir, results_dict_template):
            for xml in data:
                results_dict = deepcopy(results_dict_template)
                if not xml.endswith('.xml'):
                    continue

                xml_path = os.path.join(anno_path, xml)
                tree = ET.parse(xml_path)
                root = tree.getroot()
                # name = root.find('filename').text.split('.')[0]
                # name = '.'.join(root.find('filename').text.split('.')[:-1])
                name = xml.split('.xml')[0]
                size = root.find('size')
                if size:
                    results_dict["img_size"]["height"] = size.find(
                        'height').text
                    results_dict["img_size"]["width"] = size.find('width').text
                    results_dict["img_size"]["depth"] = size.find('depth').text
                else:
                    IMG_HEIGHT = os.getenv('img_height')
                    IMG_WIDTH = os.getenv('img_width')
                    IMG_DEPTH = os.getenv('img_depth')
                    if IMG_HEIGHT and IMG_WIDTH and IMG_DEPTH:
                        results_dict["img_size"]["height"] = IMG_HEIGHT
                        results_dict["img_size"]["width"] = IMG_WIDTH
                        results_dict["img_size"]["depth"] = IMG_DEPTH
                    else:
                        print('no size find in xml of [{}], ignore this sample.'.format(
                            xml_path))
                        continue

                anno_txt = os.path.join(gt_dir, '{}.txt'.format(name))
                with open(anno_txt, 'a') as a_f:
                    for obj in root.iter('property'):
                        t_label = obj.find('name').text
                        info = '{}'.format(t_label) + '\n'
                        a_f.writelines(info)
                        results_dict["gt"]["label"] = t_label
                info_json = os.path.join(results_dir, '{}.json'.format(name))
                with open(info_json, 'w') as f:
                    json.dump(results_dict, f)
        data = os.listdir(anno_path)
        # multiprocess
        length = len(data)
        step = int(length / g_process_num) + 1

        process_list = []
        for i in range(0, length, step):  # 开启多个子进程执行fun1函数
            p = Process(target=_process, args=(
                data[i:i+step], anno_path, gt_dir, results_dir, results_dict_template,))  # 实例化进程对象
            p.start()
            process_list.append(p)
        for p in process_list:
            p.join()


def cal_clf(methods, scoreThred, predict_path, anno_path, save_dir):
    assert methods != '', 'methods should not be ""!!'
    det_dir = os.path.join(save_dir, 'detections')
    gt_dir = os.path.join(save_dir, 'groundtruths')
    results_dir = os.path.join(save_dir, 'results')
    shutil.rmtree(det_dir, ignore_errors=True)
    shutil.rmtree(gt_dir, ignore_errors=True)
    shutil.rmtree(results_dir, ignore_errors=True)
    os.makedirs(det_dir, exist_ok=True)
    os.makedirs(gt_dir, exist_ok=True)
    os.makedirs(results_dir, exist_ok=True)

    print('start to tansfer annotation...')
    if 'coco.json' in os.listdir(anno_path):
        clf_anno_parser(anno_path,
                        gt_dir, results_dir, format='coco')
    else:
        clf_anno_parser(anno_path,
                        gt_dir, results_dir, format='voc')

    print('annotation transfered done.')
    print('start to transfer predict result...')

    data = []
    with open(predict_path, 'r') as f:
        data = f.readlines()

    def _process(lines, det_dir, results_dir):
        for line in lines:
            print(line)
            if not line:
                continue
            results_dict = {
                "img_file": "",
                "img_size": {
                    "height": "",
                    "width": "",
                    "depth": ""
                },
                "gt": {
                    "label": ""
                },
                "predicted": {
                    "predicted_label": "",
                    "scores": {}
                },
                "info": ""
            }
            det_info = line.strip().split(' ')
            name = os.path.splitext(os.path.basename(det_info[0]))[0]
            if not name:
                continue
            scores = {}
            maxScore = 0
            maxScoreLabel = ''
            for x in det_info[1].split(','):
                k, v = x.split(':')
                scores[k] = v
                results_dict["predicted"]["scores"][k] = v
                if float(v) >= maxScore:
                    maxScore = float(v)
                    maxScoreLabel = k
            results_dict["predicted"]["predicted_label"] = maxScoreLabel

            det_txt = os.path.join(det_dir, '{}.txt'.format(name))
            with open(det_txt, 'a') as d_f:
                info = '{} {}'.format(maxScoreLabel, maxScore) + '\n'
                d_f.writelines(info)
            info_json = os.path.join(results_dir, '{}.json'.format(name))
            pre_dict = ''
            try:
                with open(info_json, 'r') as f:
                    pre_dict = json.load(f)
            except Exception as e:
                print(e)
            if not pre_dict:
                pre_dict = results_dict
            else:
                pre_dict["predicted"] = results_dict["predicted"]
            pre_dict["img_file"] = det_info[0]
            if is_hikmars3:
                with open(info_json, 'w') as f:
                    json.dump(pre_dict, f)
            else:
                path = pre_dict["img_file"]
                ground_truth = [[pre_dict["gt"]["label"]]]
                predictions = [[pre_dict["predicted"]["predicted_label"],
                                pre_dict["predicted"]["scores"][pre_dict["predicted"]["predicted_label"]]]]
                EL.judge(path, ground_truth, predictions, "classify", True)
    # multiprocess
    length = len(data)
    step = int(length / g_process_num) + 1

    process_list = []
    for i in range(0, length, step):  # 开启多个子进程执行fun1函数
        p = Process(target=_process, args=(
            data[i:i+step], det_dir, results_dir,))  # 实例化进程对象
        p.start()
        process_list.append(p)
    for p in process_list:
        p.join()
    print('predict result transfered done.')

    if not is_hikmars3:
        os.system('rm -rf {}/*json'.format(results_dir))

    print('start to calucate...')
    cmd = 'python3 {}/classification/pascalvoc.py -t {} -sp {} -gt {} -det {} -methods {}'.format(
        CUR_DIR, scoreThred, save_dir, gt_dir, det_dir, methods)
    print(cmd)
    os.system(cmd)
    return


def cal_det(methods, scoreThred, IOUThred, predict_path, anno_path, save_dir):
    assert methods != '', 'methods should not be ""!!'
    det_dir = os.path.join(save_dir, 'detections')
    gt_dir = os.path.join(save_dir, 'groundtruths')
    results_dir = os.path.join(save_dir, 'results')
    shutil.rmtree(det_dir, ignore_errors=True)
    shutil.rmtree(gt_dir, ignore_errors=True)
    shutil.rmtree(results_dir, ignore_errors=True)
    os.makedirs(det_dir, exist_ok=True)
    os.makedirs(gt_dir, exist_ok=True)
    os.makedirs(results_dir, exist_ok=True)

    print('start to tansfer annotation...')
    if 'coco.json' in os.listdir(anno_path):
        det_anno_parser(anno_path,
                        gt_dir, results_dir, format='coco')
    else:
        det_anno_parser(anno_path,
                        gt_dir, results_dir, format='voc')

    print('annotation transfered done.')
    print('start to transfer predict result...')

    data = []
    with open(predict_path, 'r') as f:
        data = f.readlines()

    def _process(lines, det_dir, results_dir):
        for line in lines:
            print(line)
            if not line:
                continue
            results_dict = {
                "img_file": "",
                "img_size": {
                    "height": "",
                    "width": "",
                    "depth": ""
                },
                "gt": {
                    "object": []
                },
                "predicted": {
                    "detected_object": []
                },
                "info": ""
            }
            det_info = line.strip().split(' ')
            # name = det_info[0].split('.')[0]
            name = os.path.splitext(os.path.basename(det_info[0]))[0]
            if not name:
                continue

            det_txt = os.path.join(det_dir, '{}.txt'.format(name))

            with open(det_txt, 'a') as d_f:
                # for i in range(1, len(det_info), 6):
                index = -1
                for i in range(1, len(det_info), 5):
                    index += 1
                    t_obj = {
                        "predicted_label": "",
                        "scores": {},
                        "bbox": {
                            "xmin": 0,
                            "ymin": 0,
                            "xmax": 0,
                            "ymax": 0
                        }
                    }
                    scores_str = det_info[i]
                    scores = {}
                    maxScore = 0
                    maxScoreLabel = ''
                    for x in scores_str.split(','):
                        k, v = x.split(':')
                        scores[k] = v
                        t_obj["scores"][k] = v
                        if float(v) >= maxScore:
                            maxScore = float(v)
                            maxScoreLabel = k
                    t_obj["predicted_label"] = maxScoreLabel
                    t_obj["bbox"]["xmin"] = det_info[i + 1]
                    t_obj["bbox"]["ymin"] = det_info[i + 2]
                    t_obj["bbox"]["xmax"] = det_info[i + 3]
                    t_obj["bbox"]["ymax"] = det_info[i + 4]
                    t_obj["index"] = index
                    results_dict["predicted"]["detected_object"].append(t_obj)
                    info = '{} {} {} {}'.format(
                        maxScoreLabel, maxScore, ' '.join(
                            det_info[i + 1:i + 5]), index) + '\n'
                    d_f.writelines(info)
            info_json = os.path.join(results_dir, '{}.json'.format(name))
            pre_dict = {}
            try:
                with open(info_json, 'r') as f:
                    pre_dict = json.load(f)
            except Exception as e:
                pass
            if not pre_dict:
                pre_dict = results_dict
            else:
                pre_dict["predicted"] = results_dict["predicted"]
            pre_dict["img_file"] = det_info[0]
            if is_hikmars3:
                with open(info_json, 'w') as f:
                    json.dump(pre_dict, f)
            else:
                h, w = float(pre_dict["img_size"]["height"]), float(
                    pre_dict["img_size"]["width"])
                path = pre_dict["img_file"]
                ground_truth = []
                object_data = pre_dict["gt"]["object"]
                for obj in object_data:
                    ground_truth.append([obj["label"], float(obj["bbox"]["xmin"])/w, float(
                        obj["bbox"]["ymin"])/h, float(obj["bbox"]["xmax"])/w, float(obj["bbox"]["ymax"])/h])
                predictions = []
                pre_data = pre_dict["predicted"]["detected_object"]
                for obj in pre_data:
                    predictions.append([obj["predicted_label"], obj["scores"][obj["predicted_label"]], float(obj["bbox"]["xmin"])/w, float(
                        obj["bbox"]["ymin"])/h, float(obj["bbox"]["xmax"])/w, float(obj["bbox"]["ymax"])/h])
                EL.judge(path, ground_truth, predictions, "detect", True)

    # multiprocess
    length = len(data)
    step = int(length / g_process_num) + 1

    process_list = []
    for i in range(0, length, step):  # 开启多个子进程执行fun1函数
        p = Process(target=_process, args=(
            data[i:i+step], det_dir, results_dir,))  # 实例化进程对象
        p.start()
        process_list.append(p)
    for p in process_list:
        p.join()
    print('predict result transfered done.')

    if not is_hikmars3:
        os.system('rm -rf {}/*json'.format(results_dir))

    # delete gt data if not existed in detect data
    gt_datas = [x for x in os.listdir(gt_dir)]

    def _process(gt_datas, gt_dir, det_dir):
        for gt_file in gt_datas:
            if not os.path.isfile(os.path.join(det_dir, gt_file)):
                anno_path = os.path.join(gt_dir, gt_file)
                os.system('rm {}'.format(anno_path))

    length = len(gt_datas)
    step = int(length / g_process_num) + 1

    process_list = []
    for i in range(0, length, step):  # 开启多个子进程执行fun1函数
        p = Process(target=_process, args=(
            gt_datas[i:i+step], gt_dir, det_dir))  # 实例化进程对象
        p.start()
        process_list.append(p)
    for p in process_list:
        p.join()

    print('start to calucate...')
    cmd = 'python3 {}/detection/pascalvoc.py -st {} -it {} -sp {} -gt {} -det {} -methods {}'.format(
        CUR_DIR, scoreThred, IOUThred, save_dir, gt_dir, det_dir, methods)
    print(cmd)
    os.system(cmd)
    return


def calculate(type, methods, scoreThred, IOUThred, predict_path, anno_path, save_dir, process_num=10):
    global g_process_num
    g_process_num = process_num
    if type == 'img_classification':
        cal_clf(methods, scoreThred, predict_path, anno_path, save_dir)
    elif type == 'img_detection':
        cal_det(methods, scoreThred, IOUThred, predict_path, anno_path,
                save_dir)
        return
