
#!/bin/bash

# Get the argument from command line
if [ $# -eq 0 ]; then
    echo "Usage: $0 <folder_name>"
    exit 1
fi

argument=$1
models=$(find models/${argument} -maxdepth 1 -type d -name "*" | grep -v "^models/${argument}$" | sed "s|^models/${argument}/||")

for model in $models; do
    echo "Processing model type: $model"
    # Loop through each configuration
    echo "Training model with ${config} configuration..."
    python scripts/main.py \
        --model_dir models/${argument}/${model}/ \
        --gibbs_iterations 100 \
        --rho_scheduling_type linear \
        --rho_start 1.0 \
        --rho_end .5 \
        --initial_x measurement \
        --result_dir results/${argument}/${config} \
        --save_data \
        --initial_x measurement \
        --test_on test \
        --from_scratch \
        --checkpoints_to_keep 2000

    echo "Completed training for ${model} configuration."
    echo "----------------------------------------"
done

echo "All models trained successfully!"