#!/bin/bash
RED='\033[0;31m'
NC='\033[0m' # No Color


echo  "##---------------------------------"
echo  "GTA: Guided Transfer of Spatial Attention from Self-supervised Models"
echo  "##---------------------------------"

POSITIONAL_ARGS=()
while [[ $# -gt 0 ]]; do
  case $1 in
    -d|--dataset_name)
      DATASET="$2"
      shift # past argument
      shift # past value
      ;;
    -p|--master_port)
      PORT="$2"
      shift # past argument
      shift # past value
      ;;
    -g|--gpu_select)
      GPUS="$2"
      shift # past argument
      shift # past value
      ;;
    *|--*)
      echo "Unknown option $1"
      exit 1
      ;;
    *)
      POSITIONAL_ARGS+=("$1") # save positional arg
      shift # past argument
      ;;
  esac
done

set -- "${POSITIONAL_ARGS[@]}" # restore positional parameters
echo  "##---------------------------------"
echo -e "##Dataset Name  = ${RED}${DATASET}${NC}"
echo  "##---------------------------------"
echo  "##---------------------------------"
echo -e "##Selected GPUS  = ${RED}${GPUS}${NC}"
echo  "##---------------------------------"
echo  "##---------------------------------"
echo -e "##Selected PORT  = ${RED}${PORT}${NC}"
echo  "##---------------------------------"

# Make Baselines
iter=0
SEED=( 1222 1012 1125 )
SR=( 15 100 )
WEIGHTS=( ibot_small.pth )
NAMES=( ibot_small )
for (( s=0; s<${#SEED[@]}; s++))
do
  for (( i=0; i<${#SR[@]}; i++ ))
  do
    echo "CUDA_VISIBLE_DEVICES=${GPUS} MASTER_PORT=$((PORT+s)) python train.py \
    --fixed_RA_iter True \
    --fixed_iter 10 \
    --amp True \
    --config config/${DATASET}_template.yaml \
    --sample_rate ${SR[i]} \
    --seed ${SEED[s]} \
    --initial_checkpoint $WEIGHTS \
    --classifier_type deit_small_patch16_224_return_total_attn \
    --experiments_name ${DATASET}_${SR[i]}_guide_experiment \
    --experiments_subname ${DATASET}_${SR[i]}_${NAMES}_baseline_without_re \
    --tag ${DATASET}_${SR[i]}_${NAMES}_baseline_seed_${SEED[s]}" >> Baseline_${NAMES}_${DATASET}_${SEED[s]}.sh
  done
done

# Make Guide Experiments
iter=0
SEED=( 1222 1012 1125 )
SR=( 15 100 )
LAMB=( 0.1 1.0 10.0 100.0 )
WEIGHTS=( ibot_small.pth )
NAMES=( ibot_small )
for (( s=0; s<${#SEED[@]}; s++))
do
  for (( i=0; i<${#SR[@]}; i++ ))
  do
    for ((l=0; l<${#LAMB[@]}; l++ ))
    do
    echo "CUDA_VISIBLE_DEVICES=${GPUS} MASTER_PORT=$((PORT+s)) python train.py \
    --fixed_RA_iter True \
    --fixed_iter 10 \
    --amp True \
    --config config/${DATASET}_template.yaml \
    --guide True \
    --glambda ${LAMB[l]} \
    --sample_rate ${SR[i]} \
    --seed ${SEED[s]} \
    --initial_checkpoint $WEIGHTS \
    --classifier_type deit_small_patch16_224_return_total_attn \
    --experiments_name ${DATASET}_${SR[i]}_guide_experiment \
    --experiments_subname ${DATASET}_${SR[i]}_${NAMES}_ours_without_re \
    --tag ${DATASET}_${SR[i]}_${NAMES}_ours_glam_${LAMB[l]}_seed_${SEED[s]}" >> Guide_${NAMES}_${DATASET}_${SEED[s]}.sh
    done
  done
done