import json

FORMAT = 'det'

TASK = 'interactable'
# TASK = 'semantics'
# TASK = 'interaction'

DATASET_PATH = f'./data/coco_{FORMAT}/annotations/{TASK}.json'
TEST_613_PATH = f'../evaluation/gts/{FORMAT}/{TASK}/613.json'
TEST_GENRE_PATH = f'../evaluation/gts/{FORMAT}/{TASK}/genre.json'
TEST_CAT_PATH = f'../evaluation/gts/{FORMAT}/{TASK}/cat.json'
UNION_PATH = f'../evaluation/gts/{FORMAT}/{TASK}/union_test.json'
UNION3_PATH = f'../evaluation/gts/{FORMAT}/{TASK}/union3_test.json'


def get_union_test(fullset, testset1, testset2):
    
    union_test = {'images': [], 'categories': fullset['categories'], 'annotations': []}
    
    img_ids = []
    
    for img in testset1['images']:
        img_ids.append(img['id'])
        union_test['images'].append(img)
    for img in testset2['images']:
        if img['id'] not in img_ids:
            union_test['images'].append(img)
            img_ids.append(img['id'])
    
    for ann in fullset['annotations']:
        if ann['image_id'] in img_ids:
            union_test['annotations'].append(ann)
    
    return union_test


with open(DATASET_PATH, 'r') as f:
    dataset = json.load(f)
with open(TEST_613_PATH, 'r') as f:
    test_613 = json.load(f)
with open(TEST_GENRE_PATH, 'r') as f:
    test_genre = json.load(f)
with open(TEST_CAT_PATH, 'r') as f:
    test_cat = json.load(f)

union_test = get_union_test(dataset, test_613, test_genre)
with open(UNION_PATH, 'w') as f:
    json.dump(union_test, f, indent=4)
union3_test = get_union_test(dataset, union_test, test_cat)
with open(UNION3_PATH, 'w') as f:
    json.dump(union3_test, f, indent=4)