#!/bin/bash
# ============================================================
# dVLM-AD Waymo E2E 评估脚本
#
# 用法示例（在项目根目录执行）：
#
#   # 评估最新 checkpoint（单卡）
#   MODEL_PATH=output_models/my_run/checkpoint-1000 \
#     bash kewei_scripts/run_eval_dvlmad_waymo.sh
#
#   # 多卡并行（4 GPU）
#   MODEL_PATH=output_models/my_run/checkpoint-1000 \
#     NGPUS=4 bash kewei_scripts/run_eval_dvlmad_waymo.sh
#
#   # 快速验证（只跑前 20 条）
#   MODEL_PATH=output_models/my_run/checkpoint-1000 \
#     MAX_SAMPLES=20 bash kewei_scripts/run_eval_dvlmad_waymo.sh
# ============================================================

set -e

PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "${PROJECT_ROOT}"

# ─── 可配置参数（均可通过环境变量覆盖）─────────────────────────
MODEL_PATH="${MODEL_PATH:-}"

MODEL_PATH=output_models/dvlmad_waymo_dvlm_v2_32_lr1e-6_3361679/checkpoint-4510

BASE_MODEL="${BASE_MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}"
EVAL_JSON="${EVAL_JSON:-dvlmad_data/waymo_eval/dvlm-ad_waymo_e2e_val_cleaned.json}"
IMAGE_ROOT="${IMAGE_ROOT:-dvlmad_data}"
# Derive eval name from checkpoint path: last two path components (run_id/checkpoint-N)
_ckpt_rel="$(basename "$(dirname "${MODEL_PATH}")")/$(basename "${MODEL_PATH}")"
EVAL_NAME="${EVAL_NAME:-${_ckpt_rel}}"
OUTPUT_DIR="${OUTPUT_DIR:-eval_outputs/${EVAL_NAME}_test}"
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-1024}"
MAX_SAMPLES="${MAX_SAMPLES:--1}"
NGPUS="${NGPUS:-8}"
MIN_PIXELS="${MIN_PIXELS:-200704}"
MAX_PIXELS="${MAX_PIXELS:-200704}"

# MODEL_PATH=output_models/qwen_vl_3b_sharegpt4v_3336714/checkpoint-1897
EVAL_JSON="dvlmad_data/waymo_eval/dvlm-ad_waymo_e2e_test_cot_relative_cleaned.json"
NGPUS=8


# ─── 检查参数 ──────────────────────────────────────────────────
if [ -z "${MODEL_PATH}" ]; then
    echo "❌ 请设置 MODEL_PATH 环境变量，例如："
    echo "   MODEL_PATH=output_models/my_run/checkpoint-1000 bash $0"
    exit 1
fi

if [ ! -d "${MODEL_PATH}" ]; then
    echo "❌ MODEL_PATH 不存在: ${MODEL_PATH}"
    exit 1
fi

EVAL_JSON_ABS="${PROJECT_ROOT}/${EVAL_JSON}"
if [ ! -f "${EVAL_JSON_ABS}" ]; then
    echo "❌ 评估 JSON 不存在: ${EVAL_JSON_ABS}"
    exit 1
fi

# ─── 环境设置 ─────────────────────────────────────────────────
export PYTHONPATH="${PROJECT_ROOT}:${PROJECT_ROOT}/src:${PROJECT_ROOT}/chengyue_script:${PYTHONPATH}"
export TOKENIZERS_PARALLELISM=false
export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"
export MDM_VL_KEEP_VISION=1
export TORCHDYNAMO_DISABLE=1

echo "=========================================="
echo "dVLM-AD Waymo E2E 评估"
echo "=========================================="
echo "  模型路径      : ${MODEL_PATH}"
echo "  Base 模型     : ${BASE_MODEL}"
echo "  评估 JSON     : ${EVAL_JSON}"
echo "  图像根目录     : ${IMAGE_ROOT}"
echo "  输出目录      : ${OUTPUT_DIR}"
echo "  最大生成 Token : ${MAX_NEW_TOKENS}"
echo "  样本限制      : ${MAX_SAMPLES}  (-1 = 全部)"
echo "  GPU 数量      : ${NGPUS}"
echo "  Min/Max 像素  : ${MIN_PIXELS}/${MAX_PIXELS}"
echo "=========================================="
echo ""

mkdir -p "${OUTPUT_DIR}"

python3 kewei_scripts/eval_dvlmad_waymo.py \
    --model_path     "${MODEL_PATH}" \
    --base_model     "${BASE_MODEL}" \
    --eval_json      "${EVAL_JSON}" \
    --image_root     "${IMAGE_ROOT}" \
    --output_dir     "${OUTPUT_DIR}" \
    --max_new_tokens "${MAX_NEW_TOKENS}" \
    --max_samples    "${MAX_SAMPLES}" \
    --num_gpus       "${NGPUS}" \
    --min_pixels     "${MIN_PIXELS}" \
    --max_pixels     "${MAX_PIXELS}"

echo ""
echo "=========================================="
echo "✅ 评估完成！输出保存在: ${OUTPUT_DIR}"
echo "  predictions.json : 完整推理结果（含 CoT 文本）"
echo "  metrics.json     : 汇总指标（ADE / FDE）"
echo "=========================================="
