#!/bin/bash

LIST_OF_EXPERIMENTS=(
    "train_shannon_hyperspherical --config-name=ablation_stl10_simdex_batch_32"
    "train_shannon_hyperspherical --config-name=ablation_stl10_simdex_batch_128"
    "train_shannon_hyperspherical --config-name=ablation_stl10_simdex_beta_1.1"
    "train_shannon_hyperspherical --config-name=ablation_stl10_simdex_beta_1.25"
    "train_shannon_hyperspherical --config-name=ablation_stl10_simdex_beta_1.5"
    "train_shannon_hyperspherical --config-name=ablation_stl10_simdex_beta_1.75"
    "train_shannon_hyperspherical --config-name=ablation_stl10_simdex_beta_2"
)

# stop early if WANVDB_API_KEY is not set
if [ -z "$WANDB_API_KEY" ]; then
    echo "WANDB_API_KEY is not set. Please set it before running the script."
    exit 1
fi

image_name=simdexcmp
container_name=simdexablations2
data_path=./data

run_in_docker() {
    experiment=$1
    docker run --gpus all \
        --rm \
        --shm-size 50G \
        --network host \
        -e NVIDIA_VISIBLE_DEVICES=0 \
        -e WANDB_API_KEY \
        --name $container_name \
        -v "$(pwd)":/workspace/simdex-comparison \
        -v "$data_path":/data \
        $image_name \
        /bin/bash -c "cd /workspace/simdex-comparison && python -m pretrain.${experiment}"
}

for experiment in "${LIST_OF_EXPERIMENTS[@]}"; do
    echo "Running experiment: $experiment"
    run_in_docker "$experiment"
done