#!/bin/bash

LIST_OF_EXPERIMENTS=(
    "train_shannon_hyperspherical --config-name=stl10_shannon_hs_k1"
    "train_shannon_hyperspherical --config-name=cifar10_shannon_hs_k1"
    "train_shannon_hyperspherical --config-name=lc25000_shannon_hs_k1"
    "train_shannon_hyperspherical --config-name=vit_tiny_stl10_shannon_hs_k1"
    "train_shannon_hyperspherical --config-name=vit_tiny_lc25000_shannon_hs_k1"
    "train_shannon_hyperspherical --config-name=im1k_shannon_hs_k1"
)

# stop early if WANDB_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=simdexcmpim1k
container_name=shanon_neco_container
data_path=/root/data

run_in_docker() {
    experiment=$1
    docker run --gpus all \
        --rm \
        --shm-size 200G \
        --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