import os
import json
import shutil

# SPLIT = '613'
# SPLIT = 'genre'
SPLIT = 'cat'

OUTPUT_PATH = './coco_instances_results.json'
RESULT_PATH = f'../../evaluation/results/baselines/{SPLIT}/'
GTS_ROOT = '../../evaluation/gts'
TASKS = ['semantics', 'interactable']
FORMATS = ['det']

with open(OUTPUT_PATH, 'r') as f:
    preds = json.load(f)

for task in TASKS:
    for format in FORMATS:
        result_path = os.path.join(RESULT_PATH, format, task)
        os.makedirs(result_path, exist_ok=True)
        with open(os.path.join(GTS_ROOT, format, task, f'{SPLIT}.json'), 'r') as f:
            gts = json.load(f)
        gt_img_ids = [gt['image_id'] for gt in gts['annotations']]
        filtered_preds = []
        for pred in preds:
            if pred['image_id'] in gt_img_ids:
                filtered_preds.append(pred.copy())
        if task == 'interactable':
            for pred in filtered_preds:
                pred['category_id'] = 1
        with open(os.path.join(result_path, 'GPT4V-E2E.json'), 'w') as f:
            json.dump(filtered_preds, f)