#!/bin/bash

# custom config
DATA=...
TRAINER=AdvMaskTuning

#  =================
LAMB1=$1
LR=$2
EPOCH=$3
shift 3
DATASET_LIST=("$@")

LAYER_0=0
LAYER_1=11

for SEED in 1 2 3 ; do
#  if [[ ${SEED} == 1 ]]
#  then
#    SAVE_MODEL=True
#  else
#    SAVE_MODEL=False
#  fi

  for DATASET in "${DATASET_LIST[@]}"; do
    if [[ ${DATASET} == imagenet ]]
    then
#      CFG=vit_b32_ep10_notransform
      CFG=vit_b32_ep5_notransform
      CFG_FILE=imagenet_${CFG}.yaml
    else
#      CFG=vit_b32_ep30_notransform
      CFG=vit_b32_ep10_notransform
      CFG_FILE=${CFG}.yaml
    fi

#    for SHOTS in 16 8 4 2 1; do
    for SHOTS in 16 4 1; do

        DIR=./output/${TRAINER}_tecoa_align_ada_amp_htune/LAMB1_${LAMB1}_L${LAYER_0}_${LAYER_1}/shots_${SHOTS}/vit_b32_ep${EPOCH}_lr${LR}_notransform/${DATASET}/${SEED}
        if [ -d "$DIR" ]; then
            echo "Oops! The results exist at ${DIR} (so skip this job)"
        else
        python train.py \
        --root ${DATA} \
        --seed ${SEED} \
        --trainer ${TRAINER} \
        --dataset-config-file configs/datasets/${DATASET}.yaml \
        --config-file configs/trainers/ADVMASK/${CFG_FILE} \
        --output-dir ${DIR} \
        TEST.FINAL_MODEL "last_step" \
        DATASET.NUM_SHOTS ${SHOTS} \
        ADVMASK.LOSS_FN "align" \
        ADVMASK.LAMB1 ${LAMB1} \
        ADVMASK.PREC "amp" \
        OPTIM.LR ${LR} \
        OPTIM.NAME "sgd" \
        OPTIM.MAX_EPOCH ${EPOCH} \
        DATALOADER.TRAIN_X.BATCH_SIZE 4 \
        DATALOADER.TEST.BATCH_SIZE 200 \
        OPTIM.WEIGHT_DECAY 5e-04 \
        OPTIM.WARMUP_EPOCH 1 \
        OPTIM.WARMUP_TYPE "constant" \
        OPTIM.WARMUP_CONS_LR 1e-5 \
        ADVMASK.SAVE_MODEL ${SAVE_MODEL} \
        ADVMASK.LAYER0 ${LAYER_0} \
        ADVMASK.LAYER1 ${LAYER_1} \

        fi
    done
  done
done