#!/bin/bash

cd ..

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

DATASETS=(caltech101 dtd eurosat fgvc_aircraft food101 oxford_flowers oxford_pets stanford_cars sun397 ucf101)
CFGS=(rn50_ep50 rn50_ep100 rn50_ep100 rn50 rn50)
CTP=end  # class token position (end or middle)
NCTX=16  # number of context tokens
SHOTSALL=(1 2 4 8 16)  # number of shots (1, 2, 4, 8, 16)
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 i in 0 1 2 3 4
        do
            CFG=${CFGS[${i}]}
            SHOTS=${SHOTSALL[${i}]}
            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/RunAll/${DATASET}/${TRAINER}/${CFG}_${SHOTS}shots/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.CTX_INIT True \
                    TRAINER.COOP.N_CTX ${NCTX} \
                    TRAINER.COOP.CSC ${CSC} \
                    TRAINER.COOP.CLASS_TOKEN_POSITION ${CTP} \
                    LOSS.NAME prograd \
                    DATASET.NUM_SHOTS ${SHOTS} &
                fi
                sleep 5
                echo 'Next'
            done
        done
    done 
done 
