import os
import sys
import numpy as np
import json
import pandas as pd
from tqdm import tqdm
import argparse
import inspect

# Use original COCO API
# from pycocotools.coco import COCO
# from pycocotools.cocoeval import COCOeval
# from pycocotools.cocoeval import Params

# Ues modified ovod COCO API
from pycocotools_ovod.coco import COCO
from pycocotools_ovod.cocoeval import COCOeval
from pycocotools_ovod.cocoeval import Params

from pycocotools_ovod.semantic_matching import is_semantic_match, get_semantic_match_anns, gt_cat_match_path, reset_gt_cat_match_cache


def _process_temp_path(env_name: str, basename: str) -> str:
    override = os.environ.get(env_name)
    if override:
        return override
    tmp_dir = os.environ.get("ORIENTER_EVALUATION_TMPDIR", os.path.join(os.getcwd(), ".orienter_eval_tmp"))
    return os.path.join(tmp_dir, f"{basename}.{os.getpid()}.json")


tmp_ann_path = _process_temp_path("ORIENTER_TMP_ANN_PATH", "tmp_ann")


def _ensure_parent_dir(path: str) -> None:
    parent = os.path.dirname(path)
    if parent:
        os.makedirs(parent, exist_ok=True)


def _remove_temp_file(path: str) -> None:
    if os.path.exists(path):
        os.remove(path)


def cleanup_temp_outputs() -> None:
    _remove_temp_file(gt_cat_match_path)
    _remove_temp_file(tmp_ann_path)
    reset_gt_cat_match_cache()

def evaluate_results(cocoEval, params = None, display_summary = False):
    if params:
        cocoEval.params = params  
    print("IoU Thresholds: ",cocoEval.params.iouThrs)
    iou_index = {float(format(val, '.2f')): index for index, val in enumerate(cocoEval.params.iouThrs)} 
    
    cocoEval.evaluate()
    cocoEval.accumulate(p = params)
    if display_summary:
        cocoEval.summarize()

    precision = cocoEval.eval["precision"]
    recall = cocoEval.eval["recall"]
    scores = cocoEval.eval["scores"]

    return precision, recall, scores, iou_index


def my_format(x : float) -> str:
    formatted_number = '%.3e' % x
    parts = formatted_number.split('e')
    result = f"{parts[0]}e{int(parts[1]):01d}"
    return result


# Print final results
def cal_metrics(precision_array, recall_array, scores_array, iou_index, class_name=None):
    
    df = pd.DataFrame(columns=['class', 'IoU', 'mAP', 'F1-Score', 'Precision', 'Recall'])
        
    if not class_name:
        class_name = 'all'

    mask = precision_array == -1
    precision_array = np.ma.array(precision_array, mask=mask)

    for iou in iou_index.keys():

        map = precision_array[iou_index[iou], :, :, 0, -1].mean(1).mean()
        mprecision = precision_array[iou_index[iou], :, :, 0, -1].mean(1)
        
        max_f1 = -1
        max_f1_index = 0
        for i in range(mprecision.shape[0]):
            recall = i * 0.01
            precision = mprecision[i]
            f1 = 2 * precision * recall / (precision + recall + 1e-8)
            if f1 > max_f1:
                max_f1 = f1
                max_f1_index = i
        # map = my_format(map)
        # prec = my_format(mprecision[max_f1_index])
        # rec = my_format(max_f1_index * 0.01)
        # f1 = my_format(max_f1)
        prec = mprecision[max_f1_index]
        rec = max_f1_index * 0.01
        f1 = max_f1
        
        df.loc[len(df), df.columns] = [class_name, iou, map, f1, prec, rec]
        
    return df


def match_cats(coco_gt, coco_pred, eval_dimension):
    if os.path.exists(gt_cat_match_path):
        os.remove(gt_cat_match_path)
    reset_gt_cat_match_cache()
    dt_cats = set()
    for ann in coco_pred.dataset['annotations']:
        dt_cats.add(ann['category_id'])
    gt_cats = set()
    for cat in coco_gt.dataset['categories']:
        gt_cats.add(cat['name'])
    dt_cats = list(dt_cats)
    gt_cats = list(gt_cats)
    dt_cats.sort()
    gt_cats.sort()
    # dt_cats = gt_cats
    gt_cat_match = {gt_cat: [] for gt_cat in gt_cats}
    print('matching dt cats to gt cats...', file=sys.stderr)
    for gt_cat in tqdm(gt_cats):
        for dt_cat in dt_cats:
            if is_semantic_match(gt_cat, dt_cat, eval_dimension=eval_dimension):
                gt_cat_match[gt_cat].append(dt_cat)
        # print(gt_cat, gt_cat_match[gt_cat])
    _ensure_parent_dir(gt_cat_match_path)
    with open(gt_cat_match_path, 'w') as f:
        json.dump(gt_cat_match, f)


def do_evalutate(args):
    coco_gt = COCO(args.gt,)
    coco_pred = coco_gt.loadRes(args.dt)
    
    if 'ovod' in inspect.getfile(COCOeval):
        cocoEval = COCOeval(coco_gt, coco_pred, args.iouType, args.dimension) # modified cocoEval
    else:
        cocoEval = COCOeval(coco_gt, coco_pred, args.iouType) # original cocoEval
    
    # Load the default parameters for COCOEvaluation
    params = cocoEval.params

    ### Modify required parameters. Available params are: 
    # imgIds          - [all],
    # catIds          - [all],
    # iouThrs         - [.5:.05:.95],
    # areaRng,maxDets - [1 10 100],
    # iouType         - ['bbox'],useCats
    # eg. param.iouType = 'bbox'
    # params.iouThrs = np.linspace(.5, .9, int(np.round((.9 - .5) / .1)) + 1, endpoint=True)
    if 'ovod' in inspect.getfile(COCOeval):
        match_cats(coco_gt, coco_pred, args.dimension)

    # Evaluate the results
    precision, recall, scores, iou_index = evaluate_results(cocoEval, params, args.summary)

    df = cal_metrics(precision, recall, scores, iou_index)

    # Calculate metrics for each category
    # for cat in coco_gt.loadCats(coco_gt.getCatIds()):
    #     # Calculate the metrics
    #     params.catIds = [cat["id"]]
    #     precision, recall, scores, iou_index = evaluate_results(cocoEval, params, args.summary)
    #     class_df = cal_metrics(precision, recall, scores, iou_index, class_name=cat["name"])
    #     df = pd.concat([df, class_df], ignore_index=True)

    df.to_csv(args.log, index=False)


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']]
    _ensure_parent_dir(tmp_ann_path)
    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('-d', '--dimension', type=str, help='i / s for interactable / semantics', required=True)
    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('-i', '--iouType', type=str, default='bbox', help='iou type')
    parser.add_argument('-l', '--log', type=str, default="evaluation.log")
    parser.add_argument('-s', '--summary', action="store_true", help="Print summary of metrics")
    parser.add_argument('-n', '--name_id', action="store_true", help="Change category id to the corresponding name")
    args = parser.parse_args()
    
    # if os.path.exists(args.log):
    #     print(f'Output file already exists, skipping evaluation for {args.log}', file=sys.stderr)
    #     sys.exit(0)

    try:
        if args.name_id and 'ovod' in inspect.getfile(COCOeval):
            id2name(args)
        do_evalutate(args)
    finally:
        cleanup_temp_outputs()