import os
import sys

# Add the parent directory to sys.path to ensure imports work correctly
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

# Use absolute path for commands to ensure they run from the correct directory
root_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
os.chdir(root_dir)  # Change to the root directory before executing commands

common_command = "python -m train experiment=segmentation/segmentation_3D_amos_small.yaml ++task.loss=cross_entropy ++trainer.devices=1"

slice_thickness_values = [8, 16, 4, 2, 1, 32, 64, 128]
available_gpus = [0, 3, 4, 5, 6, 1, 2, 7]

gpu_mapping = {}
for i, thickness in enumerate(slice_thickness_values):
    gpu_mapping[thickness] = available_gpus[i % len(available_gpus)]

for thickness in slice_thickness_values:
    gpu_index = gpu_mapping[thickness]
    command = f"CUDA_VISIBLE_DEVICES={gpu_index} {common_command} ++wandb.name=seg_amos_st_{thickness} ++dataset.slice_thickness={thickness} &"
    print(command)
    os.system(command)
