#!/bin/bash

#cd ../..

# custom config
#DATA=~/efs/CoOp/data/
DATA=~/efs/imagenet/
TRAINER=CoOp

DATASET=$1
CFG=$2  # config file
CTP=$3  # class token position (end or middle)
NCTX=$4  # number of context tokens
SHOTS=$5  # number of shots (1, 2, 4, 8, 16)
CSC=$6  # class-specific context (False or True)
MODEL=$7  # imagenet_21k/${TRAINER}_${TAG}/${CFG}_unc${UNC}_${SHOTS}shots
TAG=$8

for SEED in 42
do
    TESTMODEL=best_val
    if [ ! $MODEL ]; then
      DIR=output/${DATASET}/${TRAINER}/${CFG}_${SHOTS}shots/nctx${NCTX}_csc${CSC}_ctp${CTP}/seed${SEED}
      if [ -d "$DIR" ]; then
          echo "Results are available in ${DIR}. Skip this job"
      else
          echo "Run this job and save the output to ${DIR}"
          python train.py \
          --seed ${SEED} \
          --trainer ${TRAINER} \
          --dataset-config-file configs/datasets/${DATASET}.yaml \
          --config-file configs/trainers/${TRAINER}/${CFG}.yaml \
          --root ${DATA} \
          --output-dir ${DIR} \
          TRAINER.COOP.N_CTX ${NCTX} \
          TRAINER.COOP.CSC ${CSC} \
          TRAINER.COOP.CLASS_TOKEN_POSITION ${CTP} \
          TEST.FINAL_MODEL ${TESTMODEL} \
          DATASET.NUM_SHOTS ${SHOTS}
      fi
    else
#      DIR=output/${DATASET}/${TRAINER}_init_from_${TAG}/${CFG}_${SHOTS}shots/nctx${NCTX}_csc${CSC}_ctp${CTP}/seed${SEED}
      DIR=output/${DATASET}/${TRAINER}_pretrained_name_finetuned_${TAG}/${CFG}_${SHOTS}shots/nctx${NCTX}_csc${CSC}_ctp${CTP}/seed${SEED}
      if [ -d "$DIR" ]; then
          echo "Results are available in ${DIR}. Skip this job"
      else
          echo "Run this job and save the output to ${DIR}"
          python train.py \
          --root ${DATA} \
          --seed ${SEED} \
          --trainer ${TRAINER} \
          --dataset-config-file configs/datasets/${DATASET}.yaml \
          --config-file configs/trainers/${TRAINER}/${CFG}.yaml \
          --output-dir ${DIR} \
          TRAINER.COOP.PRETRAINED_WEOGHT output/${MODEL}/nctx${NCTX}_cscFalse_ctpend/seed${SEED}/prompt_learner/model-best.pth.tar \
          TRAINER.COOP.N_CTX ${NCTX} \
          TRAINER.COOP.CSC ${CSC} \
          TRAINER.COOP.CLASS_TOKEN_POSITION ${CTP} \
          TEST.FINAL_MODEL ${TESTMODEL} \
          DATASET.NUM_SHOTS ${SHOTS}
      fi
    fi
done

# CUDA_VISIBLE_DEVICES=0 sh scripts/coop/main.sh imagenet vit_b16_ep50 end 16 16 False

# CUDA_VISIBLE_DEVICES=0 sh scripts/coop/main.sh oxford_pets vit_b16_ep100 end 4 16 False

# CUDA_VISIBLE_DEVICES=0 sh scripts/coop/main.sh oxford_pets vit_b16_ep100 end 4 16 False imagenet_21k_p/GPT_tutorial_sampled_softmax/vit_b16_ep5_unc1000_16shots imagenet_21k_p_GPT_tutorial_sampled_softmax_vit_b16_ep5_unc1000_16shots

# CUDA_VISIBLE_DEVICES=0 sh scripts/coop/main.sh oxford_pets rn50_ep100 end 4 16 False imagenet_21k_p/GPT_tutorial_sampled_softmax/rn50_ep5_unc1000_16shots imagenet_21k_p_GPT_tutorial_sampled_softmax_rn50_ep5_unc1000_16shots
# CUDA_VISIBLE_DEVICES=0 sh scripts/coop/main.sh oxford_flowers rn50_ep100_lr0.1 end 4 16 False imagenet_21k_p/GPT_tutorial_sampled_softmax/rn50_ep5_unc1000_16shots imagenet_21k_p_GPT_tutorial_sampled_softmax_rn50_ep5_unc1000_16shots
# CUDA_VISIBLE_DEVICES=0 sh scripts/coop/main.sh oxford_flowers rn50_ep100_lr0.1 end 4 16 False imagenet_21k_op/GPTH_tutorial_sampled_softmax_parent_label0.5_kl10/rn50_ep5_randaug2_unc1000_16shots imagenet_21k_op_GPTH_tutorial_sampled_softmax_parent_label0.5_kl10_rn50_ep5_randaug2_unc1000_16shots
# CUDA_VISIBLE_DEVICES=1 sh scripts/coop/main.sh oxford_flowers rn50_ep100_lr2e-4 end 4 16 False imagenet_21k_p/GPT_tutorial_sampled_softmax/rn50_ep5_unc1000_16shots imagenet_21k_p_GPT_tutorial_sampled_softmax_rn50_ep5_unc1000_16shots
