import pandas as pd
import os.path
import sys
import json
import logging
import contexttimer
import numpy as np

if len(sys.argv) != 4:
    print("Provide .tsv file name, images dir, output file name. e.g. python coco.py coco_captions_train2017.json /mnt/disks/data-1/flickr8k/coco_train.json coco_dataset_train.json")
    exit(1)

annotation_file = sys.argv[1]
images_dir = sys.argv[2]
output_file = sys.argv[3]

logging.info("Processing WIT dataset")

with contexttimer.Timer(prefix="Loading from tsv"):
    df = pd.read_csv(annotation_file, delimiter='\t')

images_dict = {}

lines = []

df = df[["caption_reference_description", "image_url"]]


df = df.replace('', np.nan)
df = df.dropna()



for index, caption_reference_description, image_url in df.itertuples():
    base_url = os.path.basename(image_url)  # extract base url
    stem, ext = os.path.splitext(base_url)  # split into stem and extension
    filename = f'{stem}.jpg'

    full_image_path =  images_dir+"/"+filename

    if os.path.isfile(full_image_path):
        lines.append(json.dumps({"image_path": full_image_path, "captions": [caption_reference_description]}))
    else:
        print(f"{full_image_path} doesn't exist")

train_lines = lines[:-9_001]
valid_lines = lines[-9_001:]

with open(output_file+"_train.json", "w") as f:
    f.write("\n".join(train_lines))

with open(output_file+"_val.json", "w") as f:
    f.write("\n".join(valid_lines))

logging.info(f"Processing Flicker WIT dataset done. {len(lines)} images processed.")