
set -eo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)"
mkdir -p workdir
# Use a specific Python interpreter if provided; avoids conda activation
PYTHON_BIN=/lustre/fsw/portfolios/nvr/users/chengyuew/anaconda3/envs/shiyi-vila/bin/python

DATASETS=${1:-${DATASETS:-MMBench_DEV_EN_V11}}
MODEL_PATH=${MODEL_PATH:-output_models/sharegpt4v_ar_baseline_6815272/checkpoint-1500}
# VLMEvalKit 目录（优先 third_party，其次 VLMEvalKit_local）
DEFAULT_VLM_DIR=""
if [ -d "${PROJECT_ROOT}/third_party/VLMEvalKit" ]; then
    DEFAULT_VLM_DIR="${PROJECT_ROOT}/third_party/VLMEvalKit"
elif [ -d "${PROJECT_ROOT}/VLMEvalKit_local" ]; then
    DEFAULT_VLM_DIR="${PROJECT_ROOT}/VLMEvalKit_local"
else
    DEFAULT_VLM_DIR="${PROJECT_ROOT}/VLMEvalKit_local"
fi
VLM_DIR=${VLM_DIR:-${DEFAULT_VLM_DIR}}
WORK_DIR=${WORK_DIR:-$(pwd)/eval_outputs/vlmeval}
DATASET_CLASS=${DATASET_CLASS:-ImageMCQDataset}
MODEL_KEY=${MODEL_KEY:-}
RUN_MODE=${RUN_MODE:-all}
USE_MDM=${USE_MDM:-0}

mkdir -p "${WORK_DIR}"

# if [ -z "${USE_MDM}" ] && [[ "${MODEL_PATH}" == *"sharegpt4v"* ]] && [[ "${MODEL_PATH}" != *"ar"* ]] ; then
#   echo "Auto-detected MDM model from path, enabling USE_MDM=1"
#   USE_MDM=1
# else
#   USE_MDM=0
# fi

# Resolve absolute path for model (use bash variable value)
if [ -d "${MODEL_PATH}" ]; then
    MODEL_PATH_ABS=$(${PYTHON_BIN} -c 'import os,sys; print(os.path.abspath(sys.argv[1]))' "${MODEL_PATH}")
else
    echo "Model path '${MODEL_PATH}' does not exist locally, assuming it is a HF Repo ID."
    MODEL_PATH_ABS="${MODEL_PATH}"
fi

CFG_PATH="${WORK_DIR}/$(date +%Y%m%d_%H%M%S)_vlmeval_config.json"

# Add script dir and dvlm models directory to PYTHONPATH
DVLM_DIR="$(cd "${SCRIPT_DIR}/../.." && pwd)"
export PYTHONPATH="${SCRIPT_DIR}:${DVLM_DIR}/models:${PYTHONPATH:-}"
export CFG_PATH MODEL_PATH_ABS DATASETS DATASET_CLASS MODEL_KEY PROCESSOR_PATH \
       USE_MDM MODEL_CLASS BLOCK_SIZE MAX_TOKENS MASK_TOKEN DEVICE

echo "Preparing VLMEvalKit config at: ${CFG_PATH}"
${PYTHON_BIN} "${SCRIPT_DIR}/generate_vlmeval_config.py"

echo "Running VLMEvalKit..."
cd "${VLM_DIR}"

# launch with torchrun for data-parallel VLMEvalKit
TORCHRUN_NPROC=$(nvidia-smi --list-gpus | wc -l | tr -d '[:space:]')
torchrun --nproc-per-node="${TORCHRUN_NPROC}" "${SCRIPT_DIR}/run_vlmeval.py" \
  --config "${CFG_PATH}" \
  --work-dir "${WORK_DIR}" \
  --mode "${RUN_MODE}" \
  --verbose

echo "Done. Results are under: ${WORK_DIR}/${MODEL_KEY}"
