#!/bin/bash

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

if [ $# -gt 6 ]
    then
        echo "The script should be invoked as ./run.sh <DATA_ROOT> <BACKBONE> <ARTIHIPPO_COMPONENT> <CUDA_VISIBLE_DEVICE> <SEQ> <IMNET_PATH>"
        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="5-datasets"
NUM_TASKS=6
SEQ=$5
BENCHMARK_CONFIG_FILE="$ROOT_DIR/configs/$BENCHMARK/$SEQ/${BACKBONE}.yaml"
DATA_ROOT=$1

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

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

CACHE_ROOT="$ARTIHIPPO_DIR/cache/$BENCHMARK"

LOG_WANDB=--log-wandb
WANDB_PROJECT=artihippo-iclr-5-datasets
WANDB_USERNAME=user
EXP_NAME=iclr-$ARTIHIPPO_COMPONENT-$SEQ-ee

for TASK_IDX in $(seq 1 `expr $NUM_TASKS - 1`)
do

    python3 $SCRIPT_DIR/search.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 --cache_data --cache_root $CACHE_ROOT  \
        --model $BACKBONE \
        --backbone_weights_path $BACKBONE_WEIGHTS_PATH --backbone_checkpoint_dir $BACKBONE_CHECKPOINT_DIR \
        --artihippo_component $ARTIHIPPO_COMPONENT \
        --amp \
        --dev-percent 10.0 --task-idx $TASK_IDX --load-single-task \
        --lr 0.001 --min-lr 1e-6 --warmup-lr 1e-4 --warmup-epochs 5 --cooldown-epochs 0 --epochs $EPOCHS --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

    python3  $SCRIPT_DIR/finetune.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 --cache_data --cache_root $CACHE_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 --load-single-task \
        --lr 0.001 --min-lr 1e-6 --warmup-lr 1e-4 --warmup-epochs 5 --cooldown-epochs 0 --epochs 30 --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

    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 --cache_data --cache_root $CACHE_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 $UNIT_TEST \
        $LOG_WANDB --wandb-project $WANDB_PROJECT --wandb-username $WANDB_USERNAME --exp-name $EXP_NAME
done
