#!/bin/bash

set -euo pipefail

BASE_MODEL="${BASE_MODEL:-llava-hf/llava-1.5-7b-hf}"
VAL_CSV="${VAL_CSV:-/home/users/ntu/cong045/scratch/testing/hallucination/CC3M-Dataset/data-qwenvl32b/bathroom_toilet_labels.csv}"
VAL_IMAGE_DIR="${VAL_IMAGE_DIR:-CC3M-Dataset/cc3m_images/train}"
OUTPUT_DIR="${OUTPUT_DIR:-./step4_outputs}"
NUM_PER_CATEGORY="${NUM_PER_CATEGORY:-50}"
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-300}"
MENTION_METHOD="${MENTION_METHOD:-both}"
JUDGE_DEVICE="${JUDGE_DEVICE:-cuda}"
JUDGE_MODEL="${JUDGE_MODEL:-Qwen/Qwen3-VL-32B-Instruct}"
MODEL_TYPE="${MODEL_TYPE:-lora}"
MODEL_PATH="${MODEL_PATH:-step3_lora_outputs/lora_adapter}"
INFERENCE_BACKEND="${INFERENCE_BACKEND:-transformers}"
VLLM_BATCH_SIZE="${VLLM_BATCH_SIZE:-64}"
VLLM_TENSOR_PARALLEL_SIZE="${VLLM_TENSOR_PARALLEL_SIZE:-1}"
VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.9}"
VLLM_MAX_MODEL_LEN="${VLLM_MAX_MODEL_LEN:-4096}"

if [ "${INFERENCE_BACKEND}" = "vllm" ] && [ -n "${CUDA_VISIBLE_DEVICES:-}" ] && [[ "${CUDA_VISIBLE_DEVICES}" == *GPU-* ]]; then
    _orig_cvd="${CUDA_VISIBLE_DEVICES}"
    declare -A _uuid_to_index=()
    while IFS=, read -r idx uuid; do
        idx="$(echo "${idx}" | xargs)"
        uuid="$(echo "${uuid}" | xargs)"
        if [ -n "${idx}" ] && [ -n "${uuid}" ]; then
            _uuid_to_index["${uuid}"]="${idx}"
        fi
    done < <(nvidia-smi --query-gpu=index,uuid --format=csv,noheader)

    IFS=',' read -r -a _requested <<< "${_orig_cvd}"
    _mapped=()
    _ok=1
    for raw in "${_requested[@]}"; do
        uuid="$(echo "${raw}" | xargs)"
        if [ -n "${_uuid_to_index[${uuid}]:-}" ]; then
            _mapped+=("${_uuid_to_index[${uuid}]}")
        else
            _ok=0
            break
        fi
    done

    if [ "${_ok}" -eq 1 ] && [ "${#_mapped[@]}" -gt 0 ]; then
        CUDA_VISIBLE_DEVICES="$(IFS=,; echo "${_mapped[*]}")"
        export CUDA_VISIBLE_DEVICES
        echo "Normalized CUDA_VISIBLE_DEVICES for vLLM: ${CUDA_VISIBLE_DEVICES}"
    else
        echo "WARNING: failed to map UUID CUDA_VISIBLE_DEVICES, keeping original: ${_orig_cvd}"
    fi
fi

if [ "${INFERENCE_BACKEND}" = "vllm" ]; then
    export VLLM_WORKER_MULTIPROC_METHOD="${VLLM_WORKER_MULTIPROC_METHOD:-spawn}"
    echo "VLLM_WORKER_MULTIPROC_METHOD=${VLLM_WORKER_MULTIPROC_METHOD}"
fi

if [ "${MODEL_TYPE}" = "delta_w" ]; then
    MODEL_ARG=(--checkpoint "${MODEL_PATH}")
else
    MODEL_ARG=(--model_dir "${MODEL_PATH}")
fi

python -m experiment.evaluation.validate \
    --model_type "${MODEL_TYPE}" \
    "${MODEL_ARG[@]}" \
    --base_model_name "${BASE_MODEL}" \
    --val_csv "${VAL_CSV}" \
    --val_image_dir "${VAL_IMAGE_DIR}" \
    --num_per_category "${NUM_PER_CATEGORY}" \
    --max_new_tokens "${MAX_NEW_TOKENS}" \
    --mention_method "${MENTION_METHOD}" \
    --judge_model "${JUDGE_MODEL}" \
    --judge_device "${JUDGE_DEVICE}" \
    --inference_backend "${INFERENCE_BACKEND}" \
    --vllm_batch_size "${VLLM_BATCH_SIZE}" \
    --vllm_tensor_parallel_size "${VLLM_TENSOR_PARALLEL_SIZE}" \
    --vllm_gpu_memory_utilization "${VLLM_GPU_MEMORY_UTILIZATION}" \
    --vllm_max_model_len "${VLLM_MAX_MODEL_LEN}" \
    --use_val_split \
    --output_dir "${OUTPUT_DIR}"