import json, os, argparse

def save_json(json_dict, fname):
    with open(fname, "w") as f:
        json.dump(json_dict, f, indent=4)

def load_json(fname):
    with open(fname, "r") as f:
        json_dict = json.load(f)
    return json_dict

def merge(eval_result_path):
    eval_result_files = [f for f in os.listdir(eval_result_path) if f.endswith('.json') and not 'merge' in f]
    merged_result = {}
    for fn in eval_result_files:
        if 'tempcompass' in fn:
            task_type = fn.replace('.json', '').replace('tempcompass_', '')
            if task_type=='multi_choice':
                task_type = 'multi-choice'
            lmms_results = load_json(f"{eval_result_path}/{fn}")
            results = {}
            for lmms_result in lmms_results['logs']:
                vid = lmms_result['doc']['video_id']
                dim = lmms_result['doc']['dim']
                if vid not in results:
                    results[vid] = {}
                if dim not in results[vid]:
                    results[vid][dim] = []
                
                result = lmms_result['avg_accuracy']
                result.pop('video_id')
                result.pop('dim')
                results[vid][dim].append(result)
            merged_result[task_type] = results
        else:
            task_type = fn.replace('.json', '')
            merged_result[task_type] = load_json(f"{eval_result_path}/{fn}")

    merge_file = f"{eval_result_path}/merged_result.json"
    save_json(merged_result, merge_file)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--eval_result_path", type=str, default="file/example_eval_results")
    args = parser.parse_args()

    merge(args.eval_result_path)