#!/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_batch_512"
    "train_shannon_hyperspherical --config-name=ablation_stl10_simdex_batch_1024"
    "train_shannon_hyperspherical --config-name=ablation_stl10_simdex_lr1e-2"
    "train_shannon_hyperspherical --config-name=ablation_stl10_simdex_lr1e-3"
    "train_shannon_hyperspherical --config-name=ablation_stl10_simdex_lr1e-4"
    "train_shannon_hyperspherical --config-name=ablation_stl10_simdex_beta_0"
    "train_shannon_hyperspherical --config-name=ablation_stl10_simdex_beta_0.01"
    "train_shannon_hyperspherical --config-name=ablation_stl10_simdex_beta_0.1"
    "train_shannon_hyperspherical --config-name=ablation_stl10_simdex_beta_0.25"
    "train_shannon_hyperspherical --config-name=ablation_stl10_simdex_beta_0.5"
    "train_shannon_hyperspherical --config-name=ablation_stl10_simdex_beta_0.75"
    "train_shannon_hyperspherical --config-name=ablation_stl10_simdex_beta_1"
)

# 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=simdexablations
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