import os

base_paths = [
    # ("/home/disk1/jianke_z/open-pi-zero/log/train/fractal_train_tp4_beta/2025-09-02_15-34_42/checkpoint",
    # "../config/eval/fractal"),
    ("/home/disk1/jianke_z/open-pi-zero/log/train/fractal_train_tp4_beta_dino_future_distill/2025-09-03_16-47_42/checkpoint",
     "../config/eval/fractal_dinodistill"),  # 这个在第三个窗口
    # ("/home/disk1/jianke_z/open-pi-zero/log/train/fractal_train_tp4_beta_siglip/2025-09-02_15-49_42/checkpoint",
    # "../config/eval/fractal_dino_siglip"), #这个在第一个窗口
    # ("/home/disk1/jianke_z/open-pi-zero/log/train/fractal_train_tp4_beta_dino/2025-09-07_00-03_42/checkpoint","../config/eval/fractal_dino" ), # 这个在第二个窗口
]
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))[-3:-1][::-1]  # only evaluate the last two checkpoints
        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 = [2, 4]
# execute_step = [4]
cuda_device = 3

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_fractal.sh {} {} {} {}".format(ckpt, step, config_name, cuda_device))
