#!/bin/bash
# ==============================================================================
# MathVista Evaluation Script
# Uses NVIDIA LLM Proxy with gpt-4o-mini as judge (Azure OpenAI format)
# ==============================================================================

set -eo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)"
mkdir -p workdir

# ==============================================================================
# NVIDIA LLM Proxy Configuration (Azure OpenAI format)
# ==============================================================================
# Azure OpenAI environment variables
export AZURE_OPENAI_ENDPOINT="https://llm-proxy.perflab.nvidia.com/"
export AZURE_OPENAI_DEPLOYMENT_NAME="gpt-4o-mini-20240718"
export OPENAI_API_VERSION="2024-12-01-preview"
export AZURE_OPENAI_API_KEY="${NSTORAGE_API_KEY:-eyJhbGciOiJIUzI1NiJ9.eyJpZCI6IjdhNGZlOGYyLTQ1NjgtNGQzYS04OWYzLTE4ZDk3MmMwMmQ2NyIsInNlY3JldCI6Iit0aTViNVJ6S2F5dUNOQmxxdFNvdEdWUmJaSXJrRE03RDZLUjMrUjZZbDQ9In0.oyiXE5cg19mBweht3C-u5BAuRVJ5-QkcgqCt5W1sGGw}"

# Use Bearer auth instead of api-key header (for NVIDIA LLM Proxy)
export AZURE_USE_BEARER=1

# ==============================================================================
# Model and Dataset Configuration
# ==============================================================================
PYTHON_BIN=/lustre/fsw/portfolios/nvr/users/chengyuew/anaconda3/envs/shiyi-vila/bin/python

# MathVista specific settings
DATASETS=${1:-${DATASETS:-MathVista_MINI}}
MODEL_PATH=${MODEL_PATH:-output_models/sharegpt4v_ar_baseline_6815272/checkpoint-1500}
# Processor path - set to original Qwen2-VL model or HF model ID
PROCESSOR_PATH=${PROCESSOR_PATH:-Qwen/Qwen2.5-VL-3B-Instruct}
# Use local VLMEvalKit with NVIDIA LLM Proxy support
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/mathvista}
DATASET_CLASS=${DATASET_CLASS:-MathVista}
MODEL_KEY=${MODEL_KEY:-}
RUN_MODE=${RUN_MODE:-all}
USE_MDM=${USE_MDM:-0}

# Judge model for MathVista evaluation (must match AZURE_OPENAI_DEPLOYMENT_NAME)
JUDGE_MODEL=${JUDGE_MODEL:-gpt-4o-mini-20240718}

mkdir -p "${WORK_DIR}"

# ==============================================================================
# Model Path Resolution
# ==============================================================================
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)_mathvista_config.json"

# ==============================================================================
# Export Environment Variables
# ==============================================================================
# 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 \
       USE_MDM MODEL_CLASS BLOCK_SIZE MAX_TOKENS MASK_TOKEN DEVICE
export PROCESSOR_PATH

# ==============================================================================
# Print Configuration
# ==============================================================================
echo "=============================================="
echo "MathVista Evaluation Configuration"
echo "=============================================="
echo "MODEL_PATH: ${MODEL_PATH_ABS}"
echo "PROCESSOR_PATH: ${PROCESSOR_PATH}"
echo "DATASETS: ${DATASETS}"
echo "DATASET_CLASS: ${DATASET_CLASS}"
echo "JUDGE_MODEL: ${JUDGE_MODEL}"
echo "WORK_DIR: ${WORK_DIR}"
echo "USE_MDM: ${USE_MDM}"
echo "AZURE_OPENAI_ENDPOINT: ${AZURE_OPENAI_ENDPOINT}"
echo "AZURE_OPENAI_DEPLOYMENT_NAME: ${AZURE_OPENAI_DEPLOYMENT_NAME}"
echo "=============================================="

# ==============================================================================
# Generate Config and Run Evaluation
# ==============================================================================
echo "Preparing VLMEvalKit config at: ${CFG_PATH}"
${PYTHON_BIN} "${SCRIPT_DIR}/generate_vlmeval_config.py"

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

# Launch with torchrun for data-parallel VLMEvalKit
TORCHRUN_NPROC=$(nvidia-smi --list-gpus | wc -l | tr -d '[:space:]')
echo "Using ${TORCHRUN_NPROC} GPUs"

torchrun --nproc-per-node="${TORCHRUN_NPROC}" "${SCRIPT_DIR}/run_vlmeval.py" \
  --config "${CFG_PATH}" \
  --work-dir "${WORK_DIR}" \
  --mode "${RUN_MODE}" \
  --judge "${JUDGE_MODEL}" \
  --verbose

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