#!/bin/bash

# Number of concurrent processes
MAX_JOBS=2

# Array of experiments to run (add your experiments here)
experiments=(
  # feedforward model with action chunking
  "python ../train/train_feedforward.py +config=feedforward config.dataset.shape=Sshape config.wandb.experiment_name=feedforward_Sshape config.architecture.model_params.output_dim=2 config.architecture.inference_params.action_chunks=1 config.architecture.inference_params.num_predicted_actions=2 config.dataset.window_size=2 config.architecture.inference_params.prediction_horizon=1"
  "python ../train/train_feedforward.py +config=feedforward config.dataset.shape=Sshape config.wandb.experiment_name=feedforward_Sshape config.architecture.model_params.output_dim=22 config.architecture.inference_params.action_chunks=11 config.architecture.inference_params.num_predicted_actions=22 config.dataset.window_size=12 config.architecture.inference_params.prediction_horizon=1"
  "python ../train/train_feedforward.py +config=feedforward config.dataset.shape=Sshape config.wandb.experiment_name=feedforward_Sshape config.architecture.model_params.output_dim=44 config.architecture.inference_params.action_chunks=22 config.architecture.inference_params.num_predicted_actions=44 config.dataset.window_size=23 config.architecture.inference_params.prediction_horizon=1"
  "python ../train/train_feedforward.py +config=feedforward config.dataset.shape=Sshape config.wandb.experiment_name=feedforward_Sshape config.architecture.model_params.output_dim=88 config.architecture.inference_params.action_chunks=44 config.architecture.inference_params.num_predicted_actions=88 config.dataset.window_size=45 config.architecture.inference_params.prediction_horizon=1"

  # our model with action chunking
  "python ../train/train_ours.py +config=ours config.dataset.shape=Sshape config.wandb.experiment_name=ours_Sshape config.architecture.model_params.output_dim=2 config.architecture.inference_params.num_predicted_actions=2 config.dataset.window_size=2 config.dataset.stride=1"
  "python ../train/train_ours.py +config=ours config.dataset.shape=Sshape config.wandb.experiment_name=ours_Sshape config.architecture.model_params.output_dim=22 config.architecture.inference_params.num_predicted_actions=22 config.dataset.window_size=12 config.dataset.stride=11"
  "python ../train/train_ours.py +config=ours config.dataset.shape=Sshape config.wandb.experiment_name=ours_Sshape config.architecture.model_params.output_dim=44 config.architecture.inference_params.num_predicted_actions=44 config.dataset.window_size=23 config.dataset.stride=22"
  "python ../train/train_ours.py +config=ours config.dataset.shape=Sshape config.wandb.experiment_name=ours_Sshape config.architecture.model_params.output_dim=88 config.architecture.inference_params.num_predicted_actions=88 config.dataset.window_size=45 config.dataset.stride=44"

)

# Function to run a command and manage the number of concurrent jobs
run_experiment() {
    local cmd="$1"
    echo "Running: $cmd"
    eval "$cmd"
}

# Loop over the experiments and run them
for experiment in "${experiments[@]}"; do
    # Wait for any job to finish if we reached the maximum concurrent jobs
    while [[ $(jobs -r -p | wc -l) -ge $MAX_JOBS ]]; do
        wait -n
    done

    # Start the experiment in the background
    run_experiment "$experiment" &
done

# Wait for any remaining jobs
wait

echo "All experiments completed."
