from collections import defaultdict
import os
import json
import pandas as pd
import argparse
from tqdm import tqdm

split = os.getenv('split', '')
suf_split = f'-{split}' if split else ''

ROUND = int(os.getenv('ROUND', 30))
CATEGORY_APP_MAP_FILE = f'./cat_apps{suf_split}.json'


def parse_imgid(img_id):
    if isinstance(img_id, int):
        img_id = str(img_id)
    return int(img_id[:-3]), int(img_id[-3:])


def cal_metrics(anno_data, interact_data):
    # Group annotations by image id
    img_annos = {}
    for anno in anno_data['annotations']:
        img_id = anno['image_id']
        if img_id not in img_annos:
            img_annos[img_id] = []
        img_annos[img_id].append(anno)

    effective_interacts_rate = 0
    effective_interacts_cnt = 0
    coverage_rate = 0
    # for sk in interact_data:
    # Add current batch interact points

    # Group interact points by image id
    img_interact = {}
    for interact in interact_data:
        img_id = interact['img_id']
        if img_id not in img_interact:
            img_interact[img_id] = []
        img_interact[img_id].append(interact)

    # Calculate effective interact rate
    for img_id in img_annos:
        if img_id not in img_interact:
            continue
        effective_interacts_rate += cnt_effective_interact(img_annos[img_id], img_interact[img_id]) / len(img_interact[img_id])
        effective_interacts_cnt += cnt_effective_interact(img_annos[img_id], img_interact[img_id])
    effective_interacts_rate /= len(img_annos)

    # Calculate coverage rate
    for img_id in img_annos:
        if img_id not in img_interact:
            continue
        coverage_rate += cnt_covered_anno(img_annos[img_id], img_interact[img_id]) / len(img_annos[img_id])
    coverage_rate /= len(img_annos)

    return effective_interacts_cnt, effective_interacts_rate, coverage_rate


def cnt_effective_interact(annos, interacts):
    effective_interacts = 0
    for interact in interacts:
        for anno in annos:
            if interact['X'] - anno['bbox'][0] >= 0 and \
               interact['Y'] - anno['bbox'][1] >= 0 and \
               interact['X'] - anno['bbox'][0] <= anno['bbox'][2] and \
               interact['Y'] - anno['bbox'][1] <= anno['bbox'][3]:
                effective_interacts += 1
                break
    return effective_interacts


def cnt_covered_anno(annos, interacts):
    coverage = 0
    # each interact point can only cover one interactable object
    t_interacts = []
    t_interacts.extend(interacts)
    for anno in annos:
        for interact in t_interacts:
            if interact['X'] - anno['bbox'][0] >= 0 and \
               interact['Y'] - anno['bbox'][1] >= 0 and \
               interact['X'] - anno['bbox'][0] <= anno['bbox'][2] and \
               interact['Y'] - anno['bbox'][1] <= anno['bbox'][3]:
                coverage += 1
                t_interacts.remove(interact)
                break
    return coverage


def main(args):
    if (os.path.exists(os.path.join(args.output_dir, 'effective_interacts_cnt.csv')) and
        os.path.exists(os.path.join(args.output_dir, 'effective_interacts_rate.csv')) and
        os.path.exists(os.path.join(args.output_dir, 'coverage_rate.csv'))):
        print("Evaluation results already exist. Loading existing results.")
        return

    with open(args.annotation, 'r') as f:
        anno_data = json.load(f)

    with open(args.category, 'r') as f:
        category_data = json.load(f)


    interact_series = args.interact_series
    interact_series = [f'{interact_series}{i}.json' for i in range(0, ROUND)]

    effective_interacts_cnt = defaultdict(list)
    effective_interacts_rate = defaultdict(list)
    coverage_rate = defaultdict(list)

    for interact_file in tqdm(interact_series, desc='Processing interact files'):

        with open(interact_file, 'r') as f:
            interact_data = json.load(f)

        cat_interact_data = {}

        for cat in category_data:
            cat_interact_data[cat] = []
            for K in interact_data:
                for interact in interact_data[K]:
                    app_id, img_id = parse_imgid(interact['img_id'])
                    if app_id in category_data[cat]:
                        cat_interact_data[cat].append(interact)


        if args.prediction:
            with open(args.prediction, 'r') as f:
                pred_data = json.load(f)

            pred_imgs = [pred['image_id'] for pred in pred_data]
            anno_data['images'] = [img for img in anno_data['images'] if img['id'] in pred_imgs]
            anno_data['annotations'] = [anno for anno in anno_data['annotations'] if anno['image_id'] in pred_imgs]


        for cat in cat_interact_data:
            if len(cat_interact_data[cat]) == 0:
                continue
            result = cal_metrics(anno_data, cat_interact_data[cat])
            effective_interacts_cnt[cat].append(result[0])
            effective_interacts_rate[cat].append(result[1])
            coverage_rate[cat].append(result[2])


    os.makedirs(args.output_dir, exist_ok=True)

    eff_int_cnt_df = pd.DataFrame(effective_interacts_cnt)
    eff_int_rate_df = pd.DataFrame(effective_interacts_rate)
    cov_rate_df = pd.DataFrame(coverage_rate)

    # Calculate average metrics for sorting
    avg_effective_interacts_rate = eff_int_rate_df.mean().sort_values(ascending=False)
    avg_coverage_rate = cov_rate_df.mean().sort_values(ascending=False)

    # Sort columns by average effective interacts rate
    eff_int_cnt_df = eff_int_cnt_df[avg_effective_interacts_rate.index]
    eff_int_rate_df = eff_int_rate_df[avg_effective_interacts_rate.index]
    cov_rate_df = cov_rate_df[avg_coverage_rate.index]

    eff_int_cnt_df.to_csv(os.path.join(args.output_dir, 'effective_interacts_cnt.csv'), index_label='Round')
    eff_int_rate_df.to_csv(os.path.join(args.output_dir, 'effective_interacts_rate.csv'), index_label='Round')
    cov_rate_df.to_csv(os.path.join(args.output_dir, 'coverage_rate.csv'), index_label='Round')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Evaluate the interact points')
    parser.add_argument('-a', '--annotation', type=str, help='Path to the annotation file')
    parser.add_argument('-is', '--interact_series', type=str, help='Path to the interact points file series')
    parser.add_argument('-o', '--output_dir', type=str, help='Path to the result directory')
    parser.add_argument('-p', '--prediction', type=str, help='Path to the prediction file')
    parser.add_argument('-c', '--category', type=str, help='Path to the app category file', default=CATEGORY_APP_MAP_FILE)
    args = parser.parse_args()
    main(args)