#!/bin/bash

saved_location="runs"
truncation_thresholds=(0.4 0.5 0.6 0.7 0.8 0.9)

# Dataset configurations
declare -A datasets
datasets=(
  ["pets"]=37
  ["flowers102"]=102
  ["cub200"]=200
  ["cifar10"]=10
  ["cifar100"]=100
)

# Model names
models=("vit_b_32")

# List of num_of_finetune values and corresponding budgets
num_of_finetune_list=(all)


declare -A budgets=(
  ["pets"]="2.30777359 5.625457764 12.80498505 29.33502197 73.08966064 209.6873779"
  ["flowers102"]="2.076477051 4.62702179 10.28045654 23.87797165 61.56833267 185.984375"
  ["cub200"]="1.862289429 4.332042694 9.741325378 22.38225937 57.53084564 173.9402161"
  ["cifar10"]="0.733501434 2.071178436 5.206390381 13.06269073 36.20851517 123.1485367"
  ["cifar100"]="0.793228149 2.237903595 5.551822662 13.65644836 37.19790649 124.9977112"
)

# Common setup
setup="A"
perplexity_pkl="--model.perplexity_pkl perplexity/perplexity_vit_b_32.pkl"
common_data_args="--data.train_workers 24 --data.val_workers 24"
common_model_args="--model.setup $setup --model.set_bn_eval True --model.use_sgd True --model.learning_rate 0.05 --model.momentum 0 --model.anneling_steps 50 --model.scheduler_interval epoch"
common_trainer_args="--trainer.max_epochs 50 --trainer.gradient_clip_val 2.0"
common_seed_args="--seed_everything 233"

# Loop through truncation thresholds
for var_index in "${!truncation_thresholds[@]}"; do
  truncation_threshold="${truncation_thresholds[var_index]}"
  echo "Running with truncation_threshold: ${truncation_threshold}"

  common_methods="--model.with_ASI True --model.truncation_threshold $truncation_threshold"
  common_args="$perplexity_pkl $common_data_args $common_model_args $common_trainer_args $common_seed_args $common_methods"
  # Loop through models and datasets
  for model in "${models[@]}"; do
    echo "Processing model: $model"

    model_config_args="--config configs/${model}_config.yaml"

    for dataset in "${!datasets[@]}"; do
      num_classes=${datasets[$dataset]}
      echo "  Processing dataset: $dataset with num_classes: $num_classes"
      specific_logger_args="--logger.save_dir ${saved_location}/setup$setup/$model/$dataset/ASI/var${truncation_threshold}/"
      specific_data_args="--data.setup $setup --data.name $dataset --data.data_dir data/$dataset"
      specific_model_args="--model.num_classes $num_classes"
      specific_args="$specific_logger_args $specific_data_args $specific_model_args"

      all_args="$model_config_args $common_args $specific_args"

      # Get budgets for the current dataset
      IFS=' ' read -r -a dataset_budgets <<< "${budgets[$dataset]}"
      for i in "${!num_of_finetune_list[@]}"; do
        num_of_finetune="${num_of_finetune_list[i]}"
        budget="${dataset_budgets[var_index]}"

        echo "    Running with num_of_finetune: $num_of_finetune, budget: $budget"
        python trainer_cls_linear.py ${all_args} --logger.exp_name ASI_var_${truncation_threshold}_l${num_of_finetune} --model.num_of_finetune $num_of_finetune --model.budget $budget
      done
    done
  done
done
