import json
import os


# # >>> CLEVR
# with open("../transformer/data/clevr/splits.json", "r") as f:
#     splits = json.load(f)

# # Change Captions
# with open("../transformer/data/clevr/change_captions.json", "r") as f:
#     change_captions = json.load(f)

# # no_change Captions
# with open("../transformer/data/clevr/no_change_captions.json", "r") as f:
#     no_change_captions = json.load(f)

# images = []
# annotations = []
# for image_id in splits["train"] + splits["val"] + splits["test"]:
#     images.append({
#         "id": f"{image_id:0>6}c",
#         "file_name": f"{image_id:0>6}c"
#     })
#     images.append({
#         "id": f"{image_id:0>6}nc",
#         "file_name": f"{image_id:0>6}nc"
#     })

# idx = 0
# for image_id in splits["train"] + splits["val"] + splits["test"]:
#     key = f"CLEVR_default_{image_id:0>6}.png"
#     change_caption = change_captions[key]
#     no_change_caption = no_change_captions[key]
    
#     for caption in change_caption:
#         annotations.append({
#             "image_id": f"{image_id:0>6}c",
#             "id": idx,
#             "caption": caption,
#         })
#         idx += 1
#     for caption in no_change_caption:
#         annotations.append({
#             "image_id": f"{image_id:0>6}nc",
#             "id": idx,
#             "caption": caption,
#         })
#         idx += 1

# with open("./clevr_data/total_change_captions_reformat.json", "wt") as f:
#     json.dump({"images": images, "annotations": annotations}, f)
# # <<<


# # >>> Spot-the-Diff
# with open("../transformer/data/spot/captions/filter_test.json", "r") as f:
#     spot_captions = json.load(f)

# with open("../transformer/data/spot/captions/filter_train.json", "r") as f:
#     spot_captions += json.load(f)

# with open("../transformer/data/spot/captions/filter_val.json", "r") as f:
#     spot_captions += json.load(f)

# images = []
# annotations = []
# idx = 0
# for item in spot_captions:
#     image_id, captions = item["img_id"], item["sentences"]
#     images.append({
#         "id": image_id,
#         "file_name": image_id
#     })
#     for caption in captions:
#         annotations.append({
#             "image_id": image_id,
#             "id": idx,
#             "caption": caption,
#         })
#         idx += 1

# with open("./spot_data/total_change_captions_reformat.json", "wt") as f:
#     json.dump({"images": images, "annotations": annotations}, f)
# # <<<


# # >>> Image-Editing-Request
# with open("../transformer/data/edit/test.json", "r") as f:
#     edit_captions = json.load(f)

# with open("../transformer/data/edit/train.json", "r") as f:
#     edit_captions += json.load(f)

# with open("../transformer/data/edit/val.json", "r") as f:
#     edit_captions += json.load(f)
    
# images = []
# annotations = []
# idx = 0
# for item in edit_captions:
#     image_id, captions = item["uid"], item["sents"]
#     images.append({
#         "id": image_id,
#         "file_name": image_id
#     })
#     for caption in captions:
#         annotations.append({
#             "image_id": image_id,
#             "id": idx,
#             "caption": caption,
#         })
#         idx += 1

# with open("./edit_data/total_change_captions_reformat.json", "wt") as f:
#     json.dump({"images": images, "annotations": annotations}, f)
# # <<<

# >>> Refined Spot-the-Diff
with open("./refined_spot_data/reformat_train.json", "r") as f:
    refined_spot_captions = json.load(f)
with open("./refined_spot_data/reformat_val.json", "r") as f:
    refined_spot_captions += json.load(f)
with open("./refined_spot_data/reformat_test.json", "r") as f:
    refined_spot_captions += json.load(f)
images = []
annotations = []
idx = 0
for item in refined_spot_captions:
    image_id, captions = item["img_id"], item["sentences"]
    images.append({
        "id": image_id,
        "file_name": image_id
    })
    for caption in captions:
        annotations.append({
            "image_id": image_id,
            "id": idx,
            "caption": caption,
        })
        idx += 1
        
with open("./refined_spot_data/total_change_captions_reformat.json", "wt") as f:
    json.dump({"images": images, "annotations": annotations}, f)
# <<<