script_folder=$HOME/scripts
salt=`head /dev/urandom | tr -dc a-z0-9 | head -c4`
# first dump the scontrol results
script_file=$script_folder/`date +'%Y-%m-%d_%H-%M-%S'`_tt.txt
echo $script_file
scontrol show job | grep "Command" > $script_file

launch_function () {
    TAG=$1
    SCRIPT=$2
    FULL_EVAL_C=$3
    EXTRA_ARGS=$4

    LA=tt
    QUEUE=mae,learnfair

    # imagenet-r
    if [ ! -f "checkpoints/TTT/imagenet_r/$TAG/model-final.pth" ]; then
        for split_label in {0..19}; do
            # The dataset_name here is intentional 
            if [ ! -f "checkpoints/TTT/imagenet_r/$TAG/$split_label/model-final.pth" ];  ~/devfair/noder.sh 1 1 1 0 72 $QUEUE $LA in-r $SCRIPT  \
                --data_path "datasets/imagenet_r" \
                --output_dir "checkpoints/TTT/imagenet_r/$TAG" \
                --dataset_name imagenet_c_split \
                $EXTRA_ARGS
            fi
        done
    fi

    # imagenet-a
    if [ ! -f "checkpoints/TTT/imagenet_a/$TAG/model-final.pth" ]; then
        for split_label in {0..19}; do
            # The dataset_name here is intentional
            if [ ! -f "checkpoints/TTT/imagenet_a/$TAG/$split_label/model-final.pth" ];  ~/devfair/noder.sh 1 1 1 0 72 $QUEUE $LA in-r $SCRIPT  \
                --data_path "datasets/imagenet_a" \
                --output_dir "checkpoints/TTT/imagenet_a/$TAG" \
                --dataset_name imagenet_c_split \
                $EXTRA_ARGS
            fi
        done
    fi

    # imagenet-c
    for DATASET in 'contrast' 'brightness' 'defocus_blur' 'elastic_transform' 'fog' 'frost' 'gaussian_noise' 'glass_blur' 'impulse_noise' 'jpeg_compression' 'motion_blur' 'pixelate' 'shot_noise' 'snow' 'zoom_blur'; do
        if [ $FULL_EVAL_C = "0" ]; then
            for tar in 5; do
                for split_label in {0..19}; do
                    if [ ! -f "checkpoints/TTT/${DATASET}_${tar}/$TAG/$split_label/model-final.pth" ]; then
                        ~/devfair/noder.sh 1 1 1 0 72 $QUEUE $LA in-c $SCRIPT  \
                            --data_path "datasets/imagenet_c/$DATASET/${tar}" \
                            --output_dir "checkpoints/TTT/${DATASET}_${tar}/$TAG/$split_label" \
                            --in_split $split_label \
                            --dataset_name imagenet_c_split \
                            $EXTRA_ARGS
                    fi
                done
            done
        else
            for tar in 1 2 3 4 5; do
                for split_label in {0..19}; do
                    if [ ! -f "checkpoints/TTT/${DATASET}_${tar}/$TAG/$split_label/model-final.pth" ]; then
                        ~/devfair/noder.sh 1 1 1 0 72 $QUEUE $LA in-c $SCRIPT  \
                            --data_path "datasets/imagenet_c/$DATASET/${tar}" \
                            --output_dir "checkpoints/TTT/${DATASET}_${tar}/$TAG/$split_label" \
                            --in_split $split_label \
                            --dataset_name imagenet_c_split \
                            $EXTRA_ARGS
                    fi
                done
            done
        fi
    done

}

###################################################
MODEL=mae_vit_large_patch16
RESUME_MODEL=checkpoints_mae/released/mae_pretrain_vit_large_full.pth
RESUME_FINETUNE=checkpoints_mae/released/mae_linear_vit_large.pth
RESUME_FINETUNE_VIT_HEAD=checkpoints/share/prob_lr1e-3_wd.2_blk12_ep20.pth
RESUME_FINETUNE_VIT_HEAD_AUG=checkpoints/share/aug_for_real_prob_lr1e-3_wd.5_dp.1_ep100.pth
RESUME_FINETUNE_VIT_HEAD_AUG_V2=checkpoints/share/aug_for_real_prob_lr1e-4_wd.3_dp.1_ep50.pth
RESUME_FINETUNE_VIT_HEAD_AUG_JOINT=checkpoints/share/joint_aug_lr1e-4_wd.05_ep100_mlc1..pth

RESUME_VIS_MODEL=checkpoints_mae/released/mae_visualize_vit_large.pth
RESUME_VIS_FINETUNE=checkpoints_mae/released/mae_vis_linear_vit_large.pth
RESUME_VIS_FINETUNE_VIT_HEAD=checkpoints/share/vis_lr1e-3_wd.2_blk12_ep20.pth

# Augs model 20 steps, 
# Please run only on imagenet-r and imagenet-a
#######################################################
TAG=joint-vit-head-sgd-aug-large_75-20-steps-encoder-dist
SHARED_ARGS=" \
--single_crop \
--batch_size 64 \
--accum_iter 2 \
--norm_pix_loss \
--mask_ratio 0.75 \
--steps_per_example 20 \
--optimizer_type 'sgd' \
--head_type 'vit_head' \
--classifier_depth 12 \
--blr 1e-2 \
--finetune_mode 'encoder' \
--model $MODEL \
--resume_model ${RESUME_FINETUNE_VIT_HEAD_AUG_JOINT} \
--resume_finetune ${RESUME_FINETUNE_VIT_HEAD_AUG_JOINT}" 
launch_function $TAG $script_file 1 "$SHARED_ARGS"

# Joint model with aug
###################################################
TAG=joint-vit-head-sgd-aug-large_75-encoder-dist
SHARED_ARGS=" \
--single_crop \
--batch_size 64 \
--accum_iter 2 \
--norm_pix_loss \
--mask_ratio 0.75 \
--steps_per_example 10 \
--optimizer_type 'sgd' \
--head_type 'vit_head' \
--classifier_depth 12 \
--blr 1e-2 \
--finetune_mode 'encoder' \
--model $MODEL \
--resume_model ${RESUME_FINETUNE_VIT_HEAD_AUG_JOINT} \
--resume_finetune ${RESUME_FINETUNE_VIT_HEAD_AUG_JOINT}" 
launch_function $TAG $script_file 1 "$SHARED_ARGS"

# TAG=joint-vit-head-sgd-aug-large_75-20-steps-encoder-dist
# SHARED_ARGS=" \
# --single_crop \
# --batch_size 64 \
# --accum_iter 2 \
# --norm_pix_loss \
# --mask_ratio 0.75 \
# --steps_per_example 20 \
# --optimizer_type 'sgd' \
# --head_type 'vit_head' \
# --classifier_depth 12 \
# --blr 1e-2 \
# --finetune_mode 'encoder' \
# --model $MODEL \
# --resume_model ${RESUME_FINETUNE_VIT_HEAD_AUG_JOINT} \
# --resume_finetune ${RESUME_FINETUNE_VIT_HEAD_AUG_JOINT}" 
# launch_function $TAG $script_file 1 "$SHARED_ARGS"
###################################################
rm -rf $script_file
echo ">>>>> $salt <<<<<"