import json
import argparse
import pandas as pd
from pycocotools.coco import COCO
import pycocotools.mask as maskUtils

gt_cat_match_path = 'tmp_gt_cat_match.json'
tmp_ann_path = 'tmp_ann.json'


def do_evaluate(args):
    coco_gt = COCO(args.gt,)
    coco_dt = coco_gt.loadRes(args.dt,)
    img_list = coco_gt.getImgIds()
    gt_masks = {}
    dt_masks = {}
    for ann in coco_gt.dataset['annotations']:
        seg = ann['segmentation']
        RLEs = maskUtils.frPyObjects(seg, 540, 960)
        RLE = maskUtils.merge(RLEs)
        if ann['image_id'] not in gt_masks:
            gt_masks[ann['image_id']] = RLE
        else:
            gt_masks[ann['image_id']] = maskUtils.merge([gt_masks[ann['image_id']], RLE])
    for ann in coco_dt.dataset['annotations']:
        seg = ann['segmentation']
        if type(seg['counts']) != str:
            RLEs = maskUtils.frPyObjects(seg, 540, 960)
            RLE = maskUtils.merge(RLEs)
        else:
            RLE = seg
        if ann['image_id'] not in dt_masks:
            dt_masks[ann['image_id']] = RLE
        else:
            dt_masks[ann['image_id']] = maskUtils.merge([dt_masks[ann['image_id']], RLE])
    precisions = {}
    recalls = {}
    for img_id in img_list:
        if img_id not in gt_masks or img_id not in dt_masks:
            precisions[img_id] = 0
            recalls[img_id] = 0
            continue
        tp_mask = maskUtils.merge([gt_masks[img_id], dt_masks[img_id]], intersect=True)
        gt_area = maskUtils.area(gt_masks[img_id])
        dt_area = maskUtils.area(dt_masks[img_id])
        tp_area = maskUtils.area(tp_mask)
        fp_area = dt_area - tp_area
        fn_area = gt_area - tp_area
        precisions[img_id] = tp_area / (tp_area + fp_area + 1e-8)
        recalls[img_id] = tp_area / (tp_area + fn_area + 1e-8)

    return precisions, recalls

def id2name(args):
    with open(args.gt, 'r') as f:
        gt = json.load(f)
    with open(args.dt, 'r') as f:
        dt = json.load(f)
    cat_id_name = {cat['id']: cat['name'] for cat in gt['categories']}
    for res in dt:
        res['category_id'] = cat_id_name[res['category_id']]
    with open(tmp_ann_path, 'w') as f:
        json.dump(dt, f)
    args.dt = tmp_ann_path


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Evaluate Metrics from the predictions and Ground Truths")
    parser.add_argument('-gt', '--gt', type=str, help='path to ground truth json', required=True)
    parser.add_argument('-dt', '--dt', type=str, help='path to detection json', required=True)
    parser.add_argument('-l', '--log', type=str, default="evaluation.log")
    parser.add_argument('-n', '--name_id', action="store_true", help="Change category id to the corresponding name")
    args = parser.parse_args()

    if args.name_id:
        id2name(args)
    precisions, recalls = do_evaluate(args)
    precision = sum(precisions.values()) / len(precisions)
    recall = sum(recalls.values()) / len(recalls)
    f1 = 2 * precision * recall / (precision + recall + 1e-8)
    result = {'precision': precision, 'recall': recall, 'f1': f1}
    pd.DataFrame(result, index=[0]).to_csv(args.log, index=False)