#!/bin/bash

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

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

# 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"
COMMON_TRAIN_CONFIGS="$ROOT_DIR/configs/defaults/common_train_configs.yaml"

# Backbone
BACKBONE=$2
ARTIHIPPO_COMPONENT=$3

# Benchmark
BENCHMARK=vdd
NUM_TASKS=10
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"

# Experiment settings
UNIT_TEST=''
UNIT_TEST_EPOCHS=2
SEED=42

LOG_WANDB=''
WANDB_PROJECT=cleanup-tests-artihippo-$BACKBONE-5
WANDB_USERNAME=user
EXP_NAME=ewc-correct-$SEED-$ARTIHIPPO_COMPONENT

EWC_LAMBDA=1e6
USE_IMP=--use-imp

python3  $SCRIPT_DIR/train_ewc.py --config $COMMON_TRAIN_CONFIGS --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 \
    --artihippo_component $ARTIHIPPO_COMPONENT \
    --amp \
    --dev-percent 0.0 --task-idx 9 --offset_task_labels \
    --lr 1e-3 --min-lr 1e-6 --warmup-lr 1e-4 --warmup-epochs 5 --cooldown-epochs 0 --epochs 20 --batch-size 512 --repeated-sampling --min-batch-size 128  \
    --unit-test-epochs $UNIT_TEST_EPOCHS $UNIT_TEST $LOG_WANDB --wandb-project $WANDB_PROJECT --wandb-username $WANDB_USERNAME --exp-name $EXP_NAME \
    --seed $SEED --ewc-lambda $EWC_LAMBDA $USE_IMP --smoothing 0.0
