#!/bin/bash

if [ $# -lt 4 ]
    then
        echo "The script should be invoked as ./evaluate.sh <DATA_ROOT> <BACKBONE> <ARTIHIPPO_COMPONENT> <CUDA_VISIBLE_DEVICE>"    
        exit 1
fi

if [ $# -gt 4 ]
    then
        echo "The script should be invoked as ./evaluate.sh <DATA_ROOT> <BACKBONE> <ARTIHIPPO_COMPONENT> <CUDA_VISIBLE_DEVICE>"
        exit 1
fi

# Setup
# https://stackoverflow.com/questions/242538/unix-shell-script-find-out-which-directory-the-script-file-resides
# Absolute path to this script
SCRIPT=$(readlink -f "$0")
SCRIPT_DIR=$(dirname "$SCRIPT")
ARTIHIPPO_DIR="$SCRIPT_DIR/../.."
TIMM_DIR="$ARTIHIPPO_DIR/external/pytorch-image-models"

export PYTHONPATH="$PYTHONPATH:$TIMM_DIR:$ARTIHIPPO_DIR"
export CUDA_VISIBLE_DEVICES=$4

ROOT_DIR="$ARTIHIPPO_DIR"

# Backbone
BACKBONE=$2
ARTIHIPPO_COMPONENT=$3

# Benchmark
BENCHMARK=vdd
BENCHMARK_CONFIG_FILE="$ROOT_DIR/configs/$BENCHMARK/${BACKBONE}.yaml"
DATA_ROOT=$1

# ImageNet
IMAGENET_ROOT=$DATA_ROOT
BACKBONE_WEIGHTS_PATH="$ARTIHIPPO_DIR/artifacts/backbone_checkpoints/$BACKBONE/imagenet12.pth.tar"
BACKBONE_CHECKPOINT_DIR="$ARTIHIPPO_DIR/artifacts/correct_checkpoints_w_mean/$BACKBONE"

UNIT_TEST=''

LOG_WANDB=--log-wandb
WANDB_PROJECT=artihippo-iclr
WANDB_USERNAME=user
EXP_NAME=iclr-$ARTIHIPPO_COMPONENT-ee-eval

TASK_IDX=0

python3 $SCRIPT_DIR/evaluate.py --root_dir $ROOT_DIR \
    --benchmark $BENCHMARK --runtime-config-file $BENCHMARK_CONFIG_FILE \
    --data_root $DATA_ROOT --imagenet_root $IMAGENET_ROOT \
    --model $BACKBONE \
    --backbone_weights_path $BACKBONE_WEIGHTS_PATH --backbone_checkpoint_dir $BACKBONE_CHECKPOINT_DIR \
    --artihippo_component $ARTIHIPPO_COMPONENT \
    --amp \
    --dev-percent 0.0 --task-idx $TASK_IDX --batch-size 1024 \
    $LOG_WANDB --wandb-project $WANDB_PROJECT --wandb-username $WANDB_USERNAME --exp-name $EXP_NAME $UNIT_TEST \
