#!/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=("swinT")

num_of_finetune_list=(all)


# Common setup
setup="A"
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 truncation_threshold in "${truncation_thresholds[@]}"; do
  echo "Running with truncation_threshold: $truncation_threshold"

  common_methods="--model.with_WSI True --model.truncation_threshold $truncation_threshold"
  common_args="$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/WSI/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"

      for i in "${!num_of_finetune_list[@]}"; do
        num_of_finetune="${num_of_finetune_list[i]}"

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