#!/bin/bash
# ==============================================================================
# Universal VLM Evaluation Script
# Supports: GQA, MMMU, TextVQA, POPE, MathVista, MMBench, and more
# 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)"

# ==============================================================================
# Benchmark Presets
# ==============================================================================
# Map benchmark names to their dataset names, classes, and judge requirements
declare -A BENCHMARK_DATASETS=(
    ["gqa"]="GQA_TestDev_Balanced"
    ["mmmu"]="MMMU_DEV_VAL"
    ["textvqa"]="TextVQA_VAL"
    ["pope"]="POPE"
    ["mathvista"]="MathVista_MINI"
    ["mmbench"]="MMBench_DEV_EN_V11"
    ["mme"]="MME"
    ["seed"]="SEEDBench_IMG"
    ["ai2d"]="AI2D_TEST"
    ["scienceqa"]="ScienceQA_VAL"
    ["realworldqa"]="RealWorldQA"
    ["ocrvqa"]="OCRVQA_TEST"
    ["docvqa"]="DocVQA_VAL"
    ["chartqa"]="ChartQA_TEST"
)

# Benchmarks that require a judge model (MCQ and Y/N types)
JUDGE_REQUIRED="mmmu pope mme mmbench seed ai2d scienceqa realworldqa"

# ==============================================================================
# Parse Arguments
# ==============================================================================
show_help() {
    echo "Usage: $0 <benchmark> [benchmark2 ...] [options]"
    echo ""
    echo "Benchmarks:"
    echo "  gqa        - GQA TestDev Balanced (VQA, exact match)"
    echo "  mmmu       - MMMU DEV_VAL (MCQ, requires judge)"
    echo "  textvqa    - TextVQA VAL (VQA, vqa_score)"
    echo "  pope       - POPE (Y/N, requires judge)"
    echo "  mathvista  - MathVista MINI (VQA, requires judge)"
    echo "  mmbench    - MMBench DEV EN V11 (MCQ, requires judge)"
    echo "  mme        - MME (Y/N, requires judge)"
    echo "  seed       - SEEDBench IMG (MCQ)"
    echo "  ai2d       - AI2D TEST (MCQ)"
    echo "  scienceqa  - ScienceQA VAL (MCQ)"
    echo "  realworldqa - RealWorldQA (MCQ)"
    echo "  ocrvqa     - OCRVQA TEST (VQA)"
    echo "  docvqa     - DocVQA VAL (VQA)"
    echo "  chartqa    - ChartQA TEST (VQA)"
    echo ""
    echo "Options (set via environment variables):"
    echo "  MODEL_PATH      - Path to model checkpoint"
    echo "  PROCESSOR_PATH  - Path to processor (default: Qwen/Qwen2.5-VL-3B-Instruct)"
    echo "  WORK_DIR        - Output directory (default: ./eval_outputs/<benchmark>)"
    echo "  USE_MDM         - Use MDM model (0 or 1, default: 0)"
    echo "  RUN_MODE        - Run mode: all, infer, eval (default: all)"
    echo ""
    echo "Examples:"
    echo "  $0 gqa"
    echo "  $0 mmmu pope textvqa"
    echo "  MODEL_PATH=/path/to/model $0 gqa textvqa"
    exit 0
}

if [ $# -eq 0 ] || [ "$1" == "-h" ] || [ "$1" == "--help" ]; then
    show_help
fi

# Collect benchmarks from arguments
BENCHMARKS="$@"

# ==============================================================================
# NVIDIA LLM Proxy Configuration (Azure OpenAI format)
# ==============================================================================
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}"
export AZURE_USE_BEARER=1

# ==============================================================================
# Model and Dataset Configuration
# ==============================================================================
PYTHON_BIN=${PYTHON_BIN:-/lustre/fsw/portfolios/nvr/users/chengyuew/anaconda3/envs/shiyi-vila/bin/python}
MODEL_PATH=${MODEL_PATH:-output_models/sharegpt4v_ar_baseline_6815272/checkpoint-1500}
PROCESSOR_PATH=${PROCESSOR_PATH:-Qwen/Qwen2.5-VL-3B-Instruct}
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}}
MODEL_KEY=${MODEL_KEY:-}
RUN_MODE=${RUN_MODE:-all}
USE_MDM=${USE_MDM:-0}
JUDGE_MODEL=${JUDGE_MODEL:-gpt-4o-mini-20240718}

# ==============================================================================
# Build Dataset List
# ==============================================================================
DATASET_LIST=""
NEEDS_JUDGE=0

for benchmark in ${BENCHMARKS}; do
    benchmark_lower=$(echo "$benchmark" | tr '[:upper:]' '[:lower:]')
    
    if [ -z "${BENCHMARK_DATASETS[$benchmark_lower]}" ]; then
        echo "Warning: Unknown benchmark '$benchmark', treating as dataset name"
        DATASET_LIST="${DATASET_LIST} ${benchmark}"
    else
        DATASET_LIST="${DATASET_LIST} ${BENCHMARK_DATASETS[$benchmark_lower]}"
    fi
    
    # Check if this benchmark requires a judge
    if [[ " $JUDGE_REQUIRED " =~ " $benchmark_lower " ]]; then
        NEEDS_JUDGE=1
    fi
done

# Trim leading space
DATASET_LIST=$(echo "$DATASET_LIST" | sed 's/^[[:space:]]*//')

if [ -z "$DATASET_LIST" ]; then
    echo "Error: No valid benchmarks specified"
    exit 1
fi

# ==============================================================================
# Set Work Directory
# ==============================================================================
BENCHMARK_NAME=$(echo "$BENCHMARKS" | tr ' ' '_')
WORK_DIR=${WORK_DIR:-$(pwd)/eval_outputs/${BENCHMARK_NAME}}
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)_vlmeval_config.json"

# ==============================================================================
# Export Environment Variables
# ==============================================================================
DVLM_DIR="$(cd "${SCRIPT_DIR}/../.." && pwd)"
export PYTHONPATH="${SCRIPT_DIR}:${DVLM_DIR}/models:${PYTHONPATH:-}"
export CFG_PATH MODEL_PATH_ABS MODEL_KEY USE_MDM MODEL_CLASS BLOCK_SIZE MAX_TOKENS MASK_TOKEN DEVICE
export PROCESSOR_PATH
export DATASETS="${DATASET_LIST}"
export AUTO_DETECT_CLASS=1  # Enable auto-detection of dataset class

# ==============================================================================
# Print Configuration
# ==============================================================================
echo "=============================================="
echo "VLM Evaluation Configuration"
echo "=============================================="
echo "BENCHMARKS: ${BENCHMARKS}"
echo "DATASETS: ${DATASET_LIST}"
echo "MODEL_PATH: ${MODEL_PATH_ABS}"
echo "PROCESSOR_PATH: ${PROCESSOR_PATH}"
echo "WORK_DIR: ${WORK_DIR}"
echo "USE_MDM: ${USE_MDM}"
echo "RUN_MODE: ${RUN_MODE}"
if [ $NEEDS_JUDGE -eq 1 ]; then
    echo "JUDGE_MODEL: ${JUDGE_MODEL}"
fi
echo "=============================================="

# ==============================================================================
# Generate Config and Run Evaluation
# ==============================================================================
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:]')
echo "Using ${TORCHRUN_NPROC} GPUs"

# Build command
CMD="torchrun --nproc-per-node=${TORCHRUN_NPROC} ${SCRIPT_DIR}/run_vlmeval.py"
CMD="${CMD} --config ${CFG_PATH}"
CMD="${CMD} --work-dir ${WORK_DIR}"
CMD="${CMD} --mode ${RUN_MODE}"
if [ $NEEDS_JUDGE -eq 1 ]; then
    CMD="${CMD} --judge ${JUDGE_MODEL}"
fi
CMD="${CMD} --verbose"

echo "Command: ${CMD}"
eval ${CMD}

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

