#!/bin/bash

cd ..

# custom config
DATA=/data1/CoOpData
TRAINERS=(ProGrad CoOp CoCoOp)

DATASETS=(caltech101 dtd eurosat fgvc_aircraft food101 oxford_flowers oxford_pets stanford_cars sun397 ucf101)
CFG=rn50_ep10 # rn50_ep50
CTP=end  # class token position (end or middle)
NCTX=4  # number of context tokens
SHOTS=4  # number of shots
CSC=False  # class-specific context (False or True)

GPULIST=(0 1 2 3)
GPUIDX=0

for TRAINER in ${TRAINERS[@]}
do
    for DATASET in ${DATASETS[@]}
    do 
        for SEED in 1 2 3 
        do
            while true 
            do 
                sleep 5
                let GPUIDX=(GPUIDX+1)%${#GPULIST[@]}
                let STATIDX=GPULIST[GPUIDX]+2
                stat=$(gpustat | sed -n ${STATIDX}'p' | grep -o beier | wc -l)
                if [ "$stat" -lt 1 ]
                then
                    break
                fi 
                echo $GPUIDX'N'
            done 

            DIR=output/base2new/train_base/${DATASET}/shots_${SHOTS}/${TRAINER}/${CFG}/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}"
                CUDA_VISIBLE_DEVICES=${GPULIST[${GPUIDX}]} 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.N_CTX ${NCTX} \
                TRAINER.COOP.CSC ${CSC} \
                TRAINER.COOP.CLASS_TOKEN_POSITION ${CTP} \
                DATASET.NUM_SHOTS ${SHOTS} \
                DATASET.SUBSAMPLE_CLASSES base \
                TRAINER.COOP.CTX_INIT True &
            fi
            sleep 5
            echo 'Next'
        done
    done 
done 