#!/bin/bash
# ==============================================================================
# 统一评测脚本
# 支持 VLMEvalKit 内置模型和自定义 dLLM checkpoint
# 同时支持 third-party benchmarks (Design2Code, HumanEval, MBPP) 自动路由
# ==============================================================================

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

# ==============================================================================
# Conda 环境 (可选)
# ==============================================================================
# Conda env — loaded from cluster_env.sh
source "${PROJECT_ROOT}/scripts/cluster_env.sh"
lmflow_activate_conda

# 限制 OpenBLAS 线程数，避免资源不足
export OMP_NUM_THREADS=${OMP_NUM_THREADS:-4}
export OPENBLAS_NUM_THREADS=${OPENBLAS_NUM_THREADS:-4}
export MKL_NUM_THREADS=${MKL_NUM_THREADS:-4}

# ==============================================================================
# 配置
# ==============================================================================
# 模型配置
# - MODEL: VLMEvalKit 内置模型名 (用于内置模型模式)
# - MODEL_PATH: 自定义模型路径 (设置后启用自定义模型模式)
MODEL=${MODEL:-Qwen2.5-VL-3B-Instruct}
MODEL_PATH=${MODEL_PATH:-}  # 空=使用内置模型, 非空=使用自定义模型

# 自定义模型专用配置
PROCESSOR_PATH=${PROCESSOR_PATH:-Qwen/Qwen2.5-VL-3B-Instruct}
USE_MDM=${USE_MDM:-0}  # 1=MDM 模式, 0=AR 模式
MODEL_KEY=${MODEL_KEY:-}  # 自定义模型名 (用于输出目录)
MODEL_CLASS=${MODEL_CLASS:-}  # 可手动指定 wrapper class (如 NemotronMinistralChatMDM)
BLOCK_SIZE=${BLOCK_SIZE:-}  # MDM block_size, 空=使用模型 config 中的 bd_size
THRESHOLD=${THRESHOLD:-0.9}  # MDM threshold, default 0.9
MAX_TOKENS=${MAX_TOKENS:-512}  # MDM max_tokens, 空=使用模型 config 中的 max_tokens
USE_SPECULATIVE=${USE_SPECULATIVE:-0}  # A/B test: 1=speculative decoding, 0=standard decoding
FORCE_LEFT_TO_RIGHT=${FORCE_LEFT_TO_RIGHT:-0}  # 1=L2R sequential decoding, 0=confidence-based parallel (default)
# 支持的 benchmarks
DEFAULT_TASKS="SEEDBench2_Plus AI2D_TEST RealWorldQA DocVQA_VAL ChartQA_TEST MathVista_MINI MMMU_DEV_VAL MMMU_Pro_V_COT MMBench_DEV_EN_V11 GQA_TestDev_Balanced TextVQA_VAL POPE ChartMimic_v2_direct_600 ChartMimic_v2_customized_600"
TASKS=${TASKS:-${DEFAULT_TASKS}}

# 输出目录 (如果设置了 MODEL_PATH 且未指定 OUTPUT_BASE，则存到 MODEL_PATH/eval_vlmeval)
OUTPUT_BASE=${OUTPUT_BASE:-}

# 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}}

# GPU 数量 (自动检测，确保是纯数字)
if [ -z "${NUM_GPUS}" ]; then
    NUM_GPUS=$(nvidia-smi --list-gpus 2>/dev/null | wc -l | tr -d '[:space:]' || echo 1)
    # 确保是有效数字
    if ! [[ "${NUM_GPUS}" =~ ^[0-9]+$ ]] || [ "${NUM_GPUS}" -eq 0 ]; then
        NUM_GPUS=1
    fi
fi

# Judge 模型 - 必须和 AZURE_OPENAI_DEPLOYMENT_NAME 一致
JUDGE_MODEL=${JUDGE_MODEL:-gpt-4o-mini-20240718}

# ==============================================================================
# NVIDIA LLM Proxy 配置 (Azure OpenAI 格式)
# ==============================================================================
export AZURE_OPENAI_ENDPOINT="https://llm-proxy.perflab.nvidia.com/"
export AZURE_OPENAI_DEPLOYMENT_NAME="${JUDGE_MODEL}"
export OPENAI_API_VERSION="2024-12-01-preview"
export AZURE_OPENAI_API_KEY="${AZURE_OPENAI_API_KEY:-eyJhbGciOiJIUzI1NiJ9.eyJpZCI6IjdhNGZlOGYyLTQ1NjgtNGQzYS04OWYzLTE4ZDk3MmMwMmQ2NyIsInNlY3JldCI6Iit0aTViNVJ6S2F5dUNOQmxxdFNvdEdWUmJaSXJrRE03RDZLUjMrUjZZbDQ9In0.oyiXE5cg19mBweht3C-u5BAuRVJ5-QkcgqCt5W1sGGw}"
export AZURE_USE_BEARER=1

# ==============================================================================
# 帮助信息
# ==============================================================================
show_help() {
    echo "VLMEvalKit 批量评测脚本"
    echo "支持 VLMEvalKit 内置模型和自定义 dLLM checkpoint"
    echo ""
    echo "Usage: $0 [options]"
    echo ""
    echo "=== 内置模型模式 (默认) ==="
    echo "  MODEL        - VLMEvalKit 内置模型名 (default: Qwen2.5-VL-3B-Instruct)"
    echo ""
    echo "=== 自定义模型模式 (设置 MODEL_PATH 启用) ==="
    echo "  MODEL_PATH     - 自定义模型 checkpoint 路径"
    echo "  PROCESSOR_PATH - Processor 路径 (default: Qwen/Qwen2.5-VL-3B-Instruct)"
    echo "  USE_MDM        - 1=MDM 模式, 0=AR 模式 (default: 0)"
    echo "  MODEL_CLASS    - 自定义 wrapper class (可选, default: 自动检测)"
    echo "  MODEL_KEY      - 自定义模型名 (用于输出目录, default: 从路径推断)"
    echo "  BLOCK_SIZE     - MDM block_size (default: 使用模型 config 中的 bd_size)"
    echo "  THRESHOLD      - MDM threshold (default: 0.9)"
    echo "  MAX_TOKENS     - MDM max_tokens (default: 使用模型 config 中的 max_tokens)"
    echo "  USE_SPECULATIVE- A/B test: 1=speculative decoding, 0=standard decoding (default: 0)"
    echo "  FORCE_LEFT_TO_RIGHT - 1=L2R sequential decoding, 0=confidence-based parallel (default)"
    echo ""
    echo "=== 通用配置 ==="
    echo "  EVAL_CONDA_ENV     - Conda env name for evaluation (default: shiyi-vila-eval)"
    echo "  EVAL_CONDA_ACTIVATE- 1=auto conda activate, 0=skip (default: 1)"
    echo "  TASKS        - 空格分隔的任务列表 (default: all benchmarks)"
    echo "  OUTPUT_BASE  - 输出目录"
    echo "                 自定义模型: MODEL_PATH/eval_vlmeval"
    echo "                 内置模型: eval_outputs/vlmeval_all"
    echo "  NUM_GPUS     - GPU 数量 (default: auto-detect)"
    echo "  JUDGE_MODEL  - GPT Judge 模型 (default: gpt-4o-mini-20240718)"
    echo ""
    echo "Supported Tasks:"
    echo "  SEEDBench2_Plus, AI2D_TEST, RealWorldQA, DocVQA_VAL, ChartQA_TEST, MathVista_MINI, MMMU_DEV_VAL, MMMU_Pro_V_COT, MMBench_DEV_EN_V11"
    echo "  GQA_TestDev_Balanced, TextVQA_VAL, POPE"
    echo "  MathVision, MathVision_MINI (math reasoning with CoT)"
    echo "  Image2Struct, Image2Struct_Latex (structure extraction)"
    echo "  ChartMimic_v2_direct_600, ChartMimic_v2_customized_600 (chart-to-code, 600 tasks)"
    echo "  ChartMimic_v2_direct, ChartMimic_v2_customized (chart-to-code, full 4800 tasks)"
    echo ""
    echo "Third-Party Benchmarks (standalone harness, auto-routed):"
    echo "  Design2Code (screenshot -> HTML, 484 test cases)"
    echo "  HumanEval   (text-only Python code completion, 164 problems)"
    echo "  MBPP        (text-only Python code generation, 500 problems)"
    echo ""
    echo "Examples:"
    echo "  # 内置模型"
    echo "  $0"
    echo "  MODEL=Qwen2.5-VL-7B-Instruct $0"
    echo ""
    echo "  # 自定义 AR checkpoint"
    echo "  MODEL_PATH=output_models/ar_sft/checkpoint-2000 $0"
    echo ""
    echo "  # 自定义 MDM checkpoint"
    echo "  MODEL_PATH=output_models/mdm_sft/checkpoint-2000 USE_MDM=1 $0"
    echo ""
    echo "  # Nemotron-Ministral MDM checkpoint (可显式指定 MODEL_CLASS)"
    echo "  MODEL_PATH=output_models/j63_mdm_ministral_2node_bd32_block_causal_ds/checkpoint-800 USE_MDM=1 MODEL_CLASS=NemotronMinistralChatMDM $0"
    exit 0
}

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

# ==============================================================================
# 判断模式: 内置模型 vs 自定义模型
# ==============================================================================
if [ -n "${MODEL_PATH}" ]; then
    USE_CUSTOM_MODEL=1
    
    # 解析 MODEL_PATH 为绝对路径
    # First, try from current directory
    if [ -d "${MODEL_PATH}" ]; then
        MODEL_PATH_ABS=$(cd "${MODEL_PATH}" && pwd)
    # Then try from PROJECT_ROOT (in case we're in a subdirectory)
    elif [ -d "${PROJECT_ROOT}/${MODEL_PATH}" ]; then
        MODEL_PATH_ABS=$(cd "${PROJECT_ROOT}/${MODEL_PATH}" && pwd)
    # If already absolute but directory doesn't exist locally (could be valid on compute node)
    elif [[ "${MODEL_PATH}" == /* ]]; then
        echo "Info: Using absolute MODEL_PATH as-is: ${MODEL_PATH}"
        MODEL_PATH_ABS="${MODEL_PATH}"
    else
        echo "ERROR: MODEL_PATH '${MODEL_PATH}' does not exist!"
        echo "Tried: ${MODEL_PATH} and ${PROJECT_ROOT}/${MODEL_PATH}"
        exit 1
    fi
    
    # 自动推断 MODEL_KEY (用于输出目录)
    if [ -z "${MODEL_KEY}" ]; then
        # 从路径中提取: output_models/xxx/checkpoint-2000 -> xxx_checkpoint-2000
        MODEL_KEY=$(echo "${MODEL_PATH}" | sed 's|.*/output_models/||' | tr '/' '_')
        [ -z "${MODEL_KEY}" ] && MODEL_KEY=$(basename "${MODEL_PATH}")
    fi
    
    # 自动检测 MDM 模式 (路径包含 mdm)
    if [ "${USE_MDM}" == "0" ] && echo "${MODEL_PATH}" | grep -qi "mdm"; then
        echo "Auto-detected MDM model from path"
        USE_MDM=1
    fi
    
    # 如果没有指定 OUTPUT_BASE，存到 model_root/eval_vlmeval/checkpoint_name
    if [ -z "${OUTPUT_BASE}" ]; then
        CKPT_NAME=$(basename "${MODEL_PATH_ABS}")
        if [[ "${CKPT_NAME}" == checkpoint-* ]]; then
            # MODEL_PATH is a checkpoint dir, store in model_root/eval_vlmeval/checkpoint_name
            MODEL_ROOT=$(dirname "${MODEL_PATH_ABS}")
            OUTPUT_BASE="${MODEL_ROOT}/eval_vlmeval/${CKPT_NAME}"
        else
            # MODEL_PATH is not a checkpoint dir, store in MODEL_PATH/eval_vlmeval
            OUTPUT_BASE="${MODEL_PATH_ABS}/eval_vlmeval"
        fi
    fi
else
    USE_CUSTOM_MODEL=0
    MODEL_KEY="${MODEL}"
    
    # 内置模型默认输出目录
    if [ -z "${OUTPUT_BASE}" ]; then
        OUTPUT_BASE="${PROJECT_ROOT}/eval_outputs/vlmeval_all"
    fi
fi

# ==============================================================================
# 打印配置
# ==============================================================================
echo "=============================================="
echo "VLMEvalKit 批量评测"
echo "=============================================="
if [ "${USE_CUSTOM_MODEL}" == "1" ]; then
    echo "MODE: 自定义模型"
    echo "MODEL_PATH: ${MODEL_PATH_ABS}"
    echo "PROCESSOR_PATH: ${PROCESSOR_PATH}"
    echo "USE_MDM: ${USE_MDM}"
    echo "MODEL_CLASS: ${MODEL_CLASS:-auto}"
    echo "BLOCK_SIZE: ${BLOCK_SIZE:-auto (from config)}"
    echo "THRESHOLD: ${THRESHOLD:-auto (from config)}"
    echo "USE_SPECULATIVE: ${USE_SPECULATIVE} (A/B test)"
    echo "FORCE_LEFT_TO_RIGHT: ${FORCE_LEFT_TO_RIGHT}"
else
    echo "MODE: 内置模型"
    echo "MODEL: ${MODEL}"
fi
echo "MODEL_KEY: ${MODEL_KEY}"
echo "TASKS: ${TASKS}"
echo "OUTPUT_BASE: ${OUTPUT_BASE}"
echo "NUM_GPUS: ${NUM_GPUS}"
echo "JUDGE_MODEL: ${JUDGE_MODEL}"
echo "VLM_DIR: ${VLM_DIR}"
echo "=============================================="

mkdir -p "${OUTPUT_BASE}"


# ==============================================================================
# 获取任务对应的 Judge 模型
# ==============================================================================
get_judge_for_task() {
    local task=$1
    case "${task}" in
        # MathVista 可以用任意 GPT judge
        MathVista*)
            echo "${JUDGE_MODEL}"
            ;;
        # ChartMimic requires GPT-4o for high-level visual judge
        ChartMimic*)
            echo "gpt-4o"
            ;;
        # MMMU/MMBench 只接受 gpt-4-0125, chatgpt-0125, exact_matching
        MMMU*|MMBench*)
            echo "gpt-4-0125"
            ;;
        # 其他任务不需要 judge
        *)
            echo ""
            ;;
    esac
}

# ==============================================================================
# Third-party benchmark runners (Design2Code, HumanEval, MBPP)
# These benchmarks have their own evaluation harnesses outside VLMEvalKit.
# ==============================================================================
EVAL_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"

is_third_party_task() {
    local task=$1
    case "${task}" in
        Design2Code|HumanEval|MBPP) return 0 ;;
        *) return 1 ;;
    esac
}

run_third_party_task() {
    local task=$1
    local task_output=$2
    local model_path_for_tp="${MODEL_PATH_ABS:-${MODEL_PATH}}"

    # PYTHONPATH for model imports
    local dvlm_dir
    dvlm_dir="$(cd "${SCRIPT_DIR}/../.." && pwd)"
    export PYTHONPATH="${dvlm_dir}/models:${PROJECT_ROOT}/src:${PYTHONPATH:-}"

    case "${task}" in
        Design2Code)
            echo "Routing to Design2Code standalone harness..."
            python3 "${EVAL_ROOT}/design2code/eval_design2code.py" \
                --model_path "${model_path_for_tp}" \
                --processor_path "${PROCESSOR_PATH}" \
                --output_dir "${task_output}" \
                --use_mdm "${USE_MDM}" \
                --max_tokens 2048 \
                ${BLOCK_SIZE:+--block_size ${BLOCK_SIZE}}
            ;;
        HumanEval)
            echo "Routing to HumanEval standalone harness..."
            python3 "${EVAL_ROOT}/code_eval/eval_humaneval_mbpp.py" \
                --model_path "${model_path_for_tp}" \
                --processor_path "${PROCESSOR_PATH}" \
                --task humaneval \
                --output_dir "${task_output}" \
                --max_tokens 512 \
                ${BLOCK_SIZE:+--block_size ${BLOCK_SIZE}}
            ;;
        MBPP)
            echo "Routing to MBPP standalone harness..."
            python3 "${EVAL_ROOT}/code_eval/eval_humaneval_mbpp.py" \
                --model_path "${model_path_for_tp}" \
                --processor_path "${PROCESSOR_PATH}" \
                --task mbpp \
                --output_dir "${task_output}" \
                --max_tokens 512 \
                ${BLOCK_SIZE:+--block_size ${BLOCK_SIZE}}
            ;;
    esac
}

# ==============================================================================
# 运行自定义模型评测 (使用 generate_vlmeval_config.py)
# ==============================================================================
run_custom_model() {
    local task=$1
    local task_output=$2
    local judge_arg=$3
    
    # 设置环境变量给 generate_vlmeval_config.py
    export CFG_PATH="${task_output}/config.json"
    export MODEL_PATH_ABS="${MODEL_PATH_ABS}"
    export PROCESSOR_PATH="${PROCESSOR_PATH}"
    export DATASETS="${task}"
    export MODEL_KEY="${MODEL_KEY}"
    export MODEL_CLASS="${MODEL_CLASS}"
    export USE_MDM="${USE_MDM}"
    export BLOCK_SIZE="${BLOCK_SIZE}"
    export MAX_TOKENS="${MAX_TOKENS}"
    export THRESHOLD="${THRESHOLD}"
    export USE_SPECULATIVE="${USE_SPECULATIVE}"
    export FORCE_LEFT_TO_RIGHT="${FORCE_LEFT_TO_RIGHT}"
    # Tokens-per-step stats (per task output)
    export TOKENS_PER_STEP_OUTPUT="${task_output}/tokens_per_step_stats.json"
    export TOKENS_PER_STEP_DUMP_EVERY="${TOKENS_PER_STEP_DUMP_EVERY:-1}"
    # 添加 PYTHONPATH
    DVLM_DIR="$(cd "${SCRIPT_DIR}/../.." && pwd)"
    export PYTHONPATH="${SCRIPT_DIR}:${DVLM_DIR}/models:${PYTHONPATH:-}"
    
    # 生成配置文件
    echo "Generating config: ${CFG_PATH}"
    python3 "${SCRIPT_DIR}/generate_vlmeval_config.py"
    
    # 运行评测 (use random master_port to avoid conflicts)
    MASTER_PORT=${MASTER_PORT:-$(python3 -c "import random; print(random.randint(20000,65000))")}
    torchrun --master_port="${MASTER_PORT}" --nproc-per-node="${NUM_GPUS}" "${SCRIPT_DIR}/run_vlmeval.py" \
        --config "${CFG_PATH}" \
        --work-dir "${task_output}" \
        --mode all \
        ${judge_arg} \
        --verbose
}

# ==============================================================================
# 运行内置模型评测 (直接用 VLMEvalKit run.py)
# ==============================================================================
run_builtin_model() {
    local task=$1
    local task_output=$2
    local judge_arg=$3
    
    torchrun --nproc-per-node="${NUM_GPUS}" run.py \
        --model "${MODEL}" \
        --data "${task}" \
        --work-dir "${task_output}" \
        ${judge_arg}
}

# ==============================================================================
# 运行评测
# ==============================================================================
cd "${VLM_DIR}"

FAILED_TASKS=""
SUCCESSFUL_TASKS=""

for TASK in ${TASKS}; do
    echo ""
    echo "=============================================="
    echo "Running: ${TASK}"
    echo "=============================================="
    
    TASK_OUTPUT="${OUTPUT_BASE}/${TASK}"
    mkdir -p "${TASK_OUTPUT}"
    
    # Route: third-party task or VLMEvalKit task
    if is_third_party_task "${TASK}"; then
        # Third-party benchmarks (Design2Code, HumanEval, MBPP)
        echo "Mode: third-party standalone harness"
        if run_third_party_task "${TASK}" "${TASK_OUTPUT}" 2>&1 | tee "${TASK_OUTPUT}/eval.log"; then
            echo "✓ ${TASK} completed successfully"
            SUCCESSFUL_TASKS="${SUCCESSFUL_TASKS} ${TASK}"
        else
            echo "✗ ${TASK} failed"
            FAILED_TASKS="${FAILED_TASKS} ${TASK}"
        fi
    else
        # VLMEvalKit tasks
        # 根据任务类型决定使用哪个 judge
        TASK_JUDGE=$(get_judge_for_task "${TASK}")
        if [ -n "${TASK_JUDGE}" ]; then
            echo "Using LLM Judge: ${TASK_JUDGE}"
            JUDGE_ARG="--judge ${TASK_JUDGE}"
        else
            echo "Using exact matching (no LLM Judge)"
            JUDGE_ARG=""
        fi
        
        # 运行评测
        if [ "${USE_CUSTOM_MODEL}" == "1" ]; then
            if run_custom_model "${TASK}" "${TASK_OUTPUT}" "${JUDGE_ARG}" 2>&1 | tee "${TASK_OUTPUT}/eval.log"; then
                echo "✓ ${TASK} completed successfully"
                SUCCESSFUL_TASKS="${SUCCESSFUL_TASKS} ${TASK}"
            else
                echo "✗ ${TASK} failed"
                FAILED_TASKS="${FAILED_TASKS} ${TASK}"
            fi
        else
            if run_builtin_model "${TASK}" "${TASK_OUTPUT}" "${JUDGE_ARG}" 2>&1 | tee "${TASK_OUTPUT}/eval.log"; then
                echo "✓ ${TASK} completed successfully"
                SUCCESSFUL_TASKS="${SUCCESSFUL_TASKS} ${TASK}"
            else
                echo "✗ ${TASK} failed"
                FAILED_TASKS="${FAILED_TASKS} ${TASK}"
            fi
        fi
    fi
done

# ==============================================================================
# 汇总结果
# ==============================================================================
echo ""
echo "=============================================="
echo "评测完成!"
echo "=============================================="
echo "成功: ${SUCCESSFUL_TASKS:-无}"
echo "失败: ${FAILED_TASKS:-无}"
echo ""
echo "结果目录: ${OUTPUT_BASE}"
echo ""

# 调用汇总脚本
if [ -f "${SCRIPT_DIR}/aggregate_vlmeval_results.py" ]; then
    echo "正在汇总结果..."
    python3 "${SCRIPT_DIR}/aggregate_vlmeval_results.py" \
        --input_dir "${OUTPUT_BASE}" \
        --model_name "${MODEL_KEY}" \
        --output "${OUTPUT_BASE}/summary.md"
    echo ""
    echo "汇总结果: ${OUTPUT_BASE}/summary.md"
    cat "${OUTPUT_BASE}/summary.md"
fi

echo "=============================================="


