import json 
import os
from natsort import natsorted
import cv2

# 增加命令行超参数
import argparse

arg_parser = argparse.ArgumentParser(description="Process some integers.")
# store类型
arg_parser.add_argument("--exoego", action='store_true', help="Use exoego dataset") # debug
arg_parser.add_argument("--start_id", type=str, default="0", help="Take ID to start with") # debug
args = arg_parser.parse_args()



take_id = "3c744ca5-c64a-4de3-8235-c2f542ac5056" # debug
if not args.exoego:
    json_path = "/scratch/yuqian_fu/egoexo_val_framelevel_newprompt_all_instruction.json"
else:
    json_path = "/scratch/yuqian_fu/ExoQuery_val_newprompt_all_instruction.json"
with open(json_path, 'r') as f:
    datas = json.load(f)

data_list = []
for data in datas:
    if data['video_name'] == take_id:
        data_list.append(data)
# 获取合适的第一帧
for data in data_list:
    if data['image'].split("/")[-1] == args.start_id + ".jpg":
        first_image = data
        break
print("first_image:", first_image["first_frame_image"])

new_img_id = 0
data_save = []
root_path = "/scratch/yuqian_fu/data_imgs" 
image_path = "/scratch/yuqian_fu/data_imgs/3c744ca5-c64a-4de3-8235-c2f542ac5056/cam03" # debug:target_image
target_images = natsorted(os.listdir(image_path))
for image in target_images:
        sample_img_path = os.path.join(image_path, image)
        sample_img_relpath = os.path.relpath(sample_img_path, root_path)
        img_target = cv2.imread(sample_img_path)
        height, width = img_target.shape[:2]
        image_info = {
            'file_name': sample_img_relpath,
            'height': height,
            'width': width,
        }
        sample = {
                'image': sample_img_relpath,
                'image_info': image_info,
                'anns': first_image['anns'],
                'first_frame_image': first_image['first_frame_image'], # 相对路径，只要最后data_root设置为"/scratch/yuqian_fu/camera_ready即可
                'first_frame_anns': first_image['first_frame_anns'],
                'new_img_id': new_img_id,
                'video_name': first_image['video_name'],
                'instruction': first_image['instruction'],
            }
        data_save.append(sample)
        new_img_id += 1
    

output_json_path = "/scratch/yuqian_fu/iccv_kitchen_ego_firstframe.json" # debug
with open(output_json_path, 'w') as f:
    json.dump(data_save, f)