import json
import random
import pandas as pd

DATASET_PATH = './data/coco_det/annotations/semantics.json'
CONTEXT_DATASET_PATH = './data/coco_context/annotations/context.json'


def parse_img_id(img_id):
    app_id = int(str(img_id)[:-3])
    img = int(str(img_id)[-3:])
    return app_id, img

with open(DATASET_PATH, 'r') as f:
    dataset = json.load(f)

cat_name_cat = {cat['name']: cat for cat in dataset['categories']}

img_ids = set([img['id'] for img in dataset['images']])

app_id_img = {}

for img in dataset['images']:
    app_id, img = parse_img_id(img['id'])
    if app_id not in app_id_img:
        app_id_img[app_id] = []
    app_id_img[app_id].append(img)

with open(CONTEXT_DATASET_PATH, 'r') as f:
    context_dataset = json.load(f)

context_cat_name = [cat['name'] for cat in context_dataset['categories'] if not cat['name'].endswith('-n')]
context_cat_id = [cat['id'] for cat in dataset['categories'] if cat['name'] in context_cat_name]

context_related_img_ids = []

for ann in dataset['annotations']:
    if ann['category_id'] in context_cat_id:
        context_related_img_ids.append(ann['image_id'])

context_related_img_ids = set(context_related_img_ids)

# Get the images that are not already in the context_related_img set
remaining_img_ids = img_ids - context_related_img_ids

# Split the remaining images into two groups with a ratio of 6:1
num_train = int(len(remaining_img_ids) * 6 / 7)
train_images = random.sample(remaining_img_ids, num_train)
val_images = remaining_img_ids - set(train_images)

print(len(train_images), len(val_images), len(context_related_img_ids))

df_data = []
df_data.extend([[img_id, 0] for img_id in train_images])
df_data.extend([[img_id, 1] for img_id in val_images])
df_data.extend([[img_id, 2] for img_id in context_related_img_ids])
df = pd.DataFrame(df_data, columns=['img_id', 'fold'])
df.to_csv('./fold_context.csv', index=False)