import os

base_paths = [
    # ("/home/disk1/jianke_z/open-pi-zero/log/train/bridge_train_tp4_beta/2025-08-24_00-06_42/checkpoint","bridge"),
    # ("/home/disk1/jianke_z/open-pi-zero/log/train/bridge_train_tp4_beta/2025-08-27_15-58_42-siglip/checkpoint","bridge_dino_siglip"),
    # ("/home/disk1/jianke_z/open-pi-zero/log/train/bridge_train_tp4_beta_dino_future_distill/2025-08-29_23-25_42/checkpoint","bridge_dinodistill"),
    # ("/home/disk1/jianke_z/open-pi-zero/log/train/bridge_train_tp4_beta_siglip_hyc_w_encoder_decoder/2025-09-11_21-00_42/checkpoint","bridge_dino_siglip_hyc"),
    # ("/home/disk1/jianke_z/open-pi-zero/log/train/bridge_train_tp4_beta_siglip_hyc_wo_enc/2025-09-13_10-56_42/checkpoint","bridge_dino_siglip_hyc_wo_enc"),
    ("/home/disk1/jianke_z/open-pi-zero/log/train/bridge_train_tp4_beta_siglip_hyc_w_encoder_decoder_fixvlm/2025-09-15_15-26_42/checkpoint",
     "bridge_dino_siglip_hyc")  # device 0
    # ("/home/disk1/jianke_z/open-pi-zero/log/train/bridge_train_tp4_beta_siglip_imgloss/2025-09-19_17-13_42/checkpoint","bridge_dino_siglip_imgloss") # device 1
]
ckpt_paths = []
for (base_path, config_name) in base_paths:
    ckpt_path = [(os.path.join(base_path, step_file), config_name)
                 for step_file in sorted(
                     os.listdir(base_path), key=lambda x: int(x.split("step")[-1].split(".pt")[0]), reverse=True)
                 if step_file.endswith(".pt")]
    ckpt_paths.extend(ckpt_path)

# ckpt_paths = [
#     # ("/home/disk1/jianke_z/open-pi-zero/log/train/bridge_train_tp4_beta/2025-08-24_00-06_42/checkpoint/step22695.pt","bridge"),
#     # ("/home/disk1/jianke_z/open-pi-zero/log/train/bridge_train_tp4_beta/2025-08-24_00-06_42/checkpoint/step21182.pt","bridge"),
#     ("/home/disk1/jianke_z/open-pi-zero/log/train/bridge_train_tp4_beta/2025-08-27_15-58_42-siglip/checkpoint/step22695.pt","bridge_dino_siglip"),
#     ]
# execute_step = [1, 2, 4]
execute_step = [4, 2]
cuda_device = 0

for i, (ckpt, config_name) in enumerate(ckpt_paths):
    for step in execute_step:
        print(f"Running evaluation for checkpoint {ckpt} with execute step {step}")
        os.system("bash slurm/eval_simpler_bridge.sh {} {} {} {}".format(ckpt, step, config_name, cuda_device))
