#!/bin/bash

# The SBATCH directives must appear before any executable line in this script.

#SBATCH --qos high2         # QOS (priority).
#SBATCH -N 1               # Number of nodes requested.
#SBATCH -n 1               # Number of tasks (i.e. processes).
#SBATCH --cpus-per-task=1  # Number of cores per task.
#SBATCH --gres=gpu:1       # Number of GPUs.
#SBATCH -t 60-00:00          # Time requested (D-HH:MM).
#SBATCH --nodelist=em6    # Uncomment if you need a specific machine.

# Uncomment this to have Slurm cd to a directory before running the script.
# You can also just run the script from the directory you want to be in.
#SBATCH -D /home/code/test_time_training/ttt_mae

# Uncomment to control the output files. By default stdout and stderr go to
# the same place, but if you use both commands below they'll be split up.
# %N is the hostname (if used, will create output(s) per node).
# %j is jobid.
##SBATCH -o slurm.%N.%j.out    # STDOUT
##SBATCH -e slurm.%N.%j.err    # STDERR

# Print some info for context.
source ~/.bashrc
conda activate taming
cd /home/test_time_training/ttt_mae_v1/

nvidia-smi
# Python will buffer output of your script unless you set this.
# If you're not using python, figure out how to turn off output
# buffering when stdout is a file, or else when watching your output
# script you'll only get updated every several lines printed.
export PYTHONUNBUFFERED=1

RESUME_MODEL='/home/test_time_training/ttt_mae_v1/models/imagenet/vit_head_2_layers/mae_pretrain_vit_large_full.pth'
# RESUME_FINETUNE='/home/test_time_training/ttt_mae_v1/models/imagenet/vit_head_2_layers/checkpoint-19.pth'
# RESUME_FINETUNE='/home/test_time_training/ttt_mae_v1/models/imagenet/visualize/vis_lr1e-3_wd.2_blk12_ep20.pth'
RESUME_FINETUNE='/home/test_time_training/ttt_mae_v1/models/imagenet/vit_head_with_prob_with_aug/aug_for_real_prob_lr1e-3_wd.5_dp.1_ep100.pth'
# RESUME_MODEL='/home/test_time_training/ttt_mae_v1/models/imagenet/visualize/mae_visualize_vit_large.pth'
OUTPUT_DIR_BASE='/home/test_time_training/ttt_mae/output_dir/vit_head_with_aug/no_test_time'
# First let's run the original dataset:



DATA_PATH_BASE='/scratch/data/imagenet_c/'
# #  'saturate' 'speckle_noise'  'gaussian_blur' 'spatter'
for DATASET in  'snow' 'brightness' 'fog' 'frost' 'contrast' 'glass_blur' 'motion_blur'  'shot_noise' 'defocus_blur' 'impulse_noise' 'pixelate'  'zoom_blur' 'elastic_transform' 'gaussian_noise' 'jpeg_compression'
do 
       CUDA_VISIBLE_DEVICES=5 python test_without_adaptation.py \
        --data_path "$DATA_PATH_BASE/$DATASET/5" \
        --model mae_vit_large_patch16 \
        --input_size 224 \
        --resume_model ${RESUME_MODEL} \
        --resume_finetune ${RESUME_FINETUNE} \
        --output_dir "$OUTPUT_DIR_BASE/$DATASET/base" \
        --classifier_depth 12 \
        --head_type "vit_head" 
done

# DATA_PATH='/shared/group/ilsvrc/val'
# CUDA_VISIBLE_DEVICES=3 python test_without_adaptation.py \
#         --data_path "$DATA_PATH" \
#         --model mae_vit_large_patch16 \
#         --input_size 224 \
#         --resume_model ${RESUME_MODEL} \
#         --resume_finetune ${RESUME_FINETUNE} \
#         --output_dir "$OUTPUT_DIR_BASE/original/base" \
#         --classifier_depth 12 \
#         --head_type "vit_head" 


# Print completion time.
date