import os
import random
import shutil
import json
from env import DATA_DIR, ROOT_DIR

if not os.path.exists(f"{DATA_DIR}/coco/annotations/coco_karpathy_train_full.json"):
    shutil.copy(f"{DATA_DIR}/coco/annotations/coco_karpathy_train.json", f"{DATA_DIR}/coco/annotations/coco_karpathy_train_full.json")

if not os.path.exists(f"{DATA_DIR}/coco/annotations/vqa_train_full.json"):
    shutil.copy(f"{DATA_DIR}/coco/annotations/vqa_train.json", f"{DATA_DIR}/coco/annotations/vqa_train_full.json")

with open(f"{DATA_DIR}/coco/annotations/coco_karpathy_train_full.json", "r") as f:
    data = json.load(f)

random.seed(42)
random.shuffle(data)
poisoned_sample = data[:20000]
benign_sample = data[20000:]

for sample in poisoned_sample:
    sample['true_caption'] = sample['caption']

with open(f"{DATA_DIR}/coco/annotations/caption_poisoned_placeholder.json", "w") as f:
    json.dump(poisoned_sample, f, indent=4)

with open(f"{DATA_DIR}/coco/annotations/caption_clean.json", "w") as f:
    json.dump(benign_sample, f, indent=4)
