#!/bin/bash

LIST_OF_EXPERIMENTS=(
    "train_dino --config-name=cifar10_dino"
    "train_byol --config-name=cifar10_byol"
    "train_barlowtwins --config-name=cifar10_barlowtwins"
    "train_simclr --config-name=cifar10_simclr"
    "train_supervised --config-name=cifar10_supervised"
    "train_swav --config-name=cifar10_swav"
    "train_vicreg --config-name=cifar10_vicreg"
)

# 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=simdexcmp
data_path=./data

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