#!/usr/bin/env bash
# =============================================================================
# Evaluate LoRA hallucination: run the trained model on the bathroom/toilet
# dataset and measure how often it still hallucinates a toilet.
#
# Mirrors run_build_caption_targets.sh but uses the LoRA-trained model.
#
# Two stages:
#   Stage 1:   Run LoRA LLaVA on all images -> raw captions
#   Stage 1.5: Regex coarse filter + LLM judge to confirm toilet mentions
#
# Data is loaded from HuggingFace (pbcong/bathroom-toilet) by default.
# Set CSV_PATH + IMAGE_DIR to use local files instead.
#
# Usage:
#   # Full pipeline (inference + LLM judge)
#   LORA_DIR=step3_lora_v5_outputs/run_20260317_000000/lora_adapter \
#       bash experiment/scripts/data/run_eval_lora_hallucination.sh
#
#   # With comparison against original model captions
#   LORA_DIR=step3_lora_v5_outputs/run_20260317_000000/lora_adapter \
#   ORIGINAL_TARGETS=experiment/data/caption_targets.json \
#       bash experiment/scripts/data/run_eval_lora_hallucination.sh
#
#   # Inference only (no LLM judge)
#   INFERENCE_ONLY=true \
#   LORA_DIR=step3_lora_v5_outputs/run_20260317_000000/lora_adapter \
#       bash experiment/scripts/data/run_eval_lora_hallucination.sh
#
#   # Judge an existing result file (stage 1.5 only)
#   JUDGE_EXISTING=experiment/data/lora_hallucination_results.json \
#       bash experiment/scripts/data/run_eval_lora_hallucination.sh
# =============================================================================

set -euo pipefail

PROJECT_ROOT="$(cd "$(dirname "$0")/../../.." && pwd)"
export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}"

export TORCHINDUCTOR_CACHE_DIR="${HOME}/scratch/.cache/torchinductor"
export TRITON_CACHE_DIR="${HOME}/scratch/.cache/triton"

# =============================================================================
# Paths & config
# =============================================================================
# HuggingFace dataset (default — no local files needed)
DATASET_ID="${DATASET_ID:-pbcong/bathroom-toilet}"
# Legacy local paths (set both to use local CSV + images)
CSV_PATH="${CSV_PATH:-}"
IMAGE_DIR="${IMAGE_DIR:-}"

# GPU selection (e.g. CUDA_VISIBLE_DEVICES=0,1)
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-}"

# Required: path to trained LoRA adapter directory
LORA_DIR="${LORA_DIR:-pbcong/bathroom-toilet-300}"

# Optional: path to original caption_targets.json for before/after comparison
ORIGINAL_TARGETS="${ORIGINAL_TARGETS:-experiment/data/caption_targets.json}"

OUTPUT="${OUTPUT:-experiment/data/lora_hallucination_results.json}"
BASE_MODEL="${BASE_MODEL:-llava-hf/llava-1.5-7b-hf}"
JUDGE_MODEL="${JUDGE_MODEL:-Qwen/Qwen3-8B}"
INFERENCE_ONLY="${INFERENCE_ONLY:-false}"
SKIP_JUDGE="${SKIP_JUDGE:-false}"
JUDGE_EXISTING="${JUDGE_EXISTING:-}"
BATCH_SIZE="${BATCH_SIZE:-8}"
NUM_GPUS="${NUM_GPUS:-1}"
JUDGE_BATCH_SIZE="${JUDGE_BATCH_SIZE:-64}"
JUDGE_GPU_MEMORY="${JUDGE_GPU_MEMORY:-0.85}"
JUDGE_NUM_GPUS="${JUDGE_NUM_GPUS:-1}"

echo "=========================================="
echo "Evaluate LoRA Hallucination"
echo "=========================================="
echo "Data:"
if [ -n "${CSV_PATH}" ] && [ -n "${IMAGE_DIR}" ]; then
    echo "  CSV:              ${CSV_PATH}"
    echo "  Image dir:        ${IMAGE_DIR}"
else
    echo "  Dataset:          ${DATASET_ID}"
fi
echo "  GPUs:             ${CUDA_VISIBLE_DEVICES:-all}"
echo "  Output:           ${OUTPUT}"
echo "Model:"
echo "  Base model:       ${BASE_MODEL}"
echo "  LoRA dir:         ${LORA_DIR:-<not set>}"
echo "  Judge model:      ${JUDGE_MODEL} (vLLM, DDP ${JUDGE_NUM_GPUS} GPU(s))"
echo "Pipeline:"
echo "  Inference only:   ${INFERENCE_ONLY}"
echo "  Skip judge:       ${SKIP_JUDGE}"
echo "  Judge existing:   ${JUDGE_EXISTING:-none}"
if [ -f "${ORIGINAL_TARGETS}" ]; then
    echo "  Original targets: ${ORIGINAL_TARGETS} (comparison enabled)"
else
    echo "  Original targets: not found (no comparison)"
fi
echo "=========================================="

# ---- Mode: judge existing file only ----
if [ -n "${JUDGE_EXISTING}" ]; then
    echo ""
    echo ">>> Stage 1.5: LLM judge of ${JUDGE_EXISTING}"
    python -m experiment.data.eval_lora_hallucination \
        --judge_only "${JUDGE_EXISTING}" \
        --judge_model "${JUDGE_MODEL}" \
        --judge_batch_size "${JUDGE_BATCH_SIZE}" \
        --judge_gpu_memory "${JUDGE_GPU_MEMORY}" \
        --judge_num_gpus "${JUDGE_NUM_GPUS}"

else
    # ---- Full pipeline ----
    # Build data args
    DATA_ARGS=()
    if [ -n "${CSV_PATH}" ] && [ -n "${IMAGE_DIR}" ]; then
        DATA_ARGS+=(--csv "${CSV_PATH}" --image_dir "${IMAGE_DIR}")
    else
        DATA_ARGS+=(--dataset_id "${DATASET_ID}")
    fi

    # Optional original targets for comparison
    ORIG_ARGS=()
    if [ -f "${ORIGINAL_TARGETS}" ]; then
        ORIG_ARGS+=(--original_targets "${ORIGINAL_TARGETS}")
    fi

    CMD=(
        python -m experiment.data.eval_lora_hallucination
        "${DATA_ARGS[@]}"
        --output "${OUTPUT}"
        --base_model "${BASE_MODEL}"
        --lora_dir "${LORA_DIR}"
        --judge_model "${JUDGE_MODEL}"
        --batch_size "${BATCH_SIZE}"
        --num_gpus "${NUM_GPUS}"
        --judge_batch_size "${JUDGE_BATCH_SIZE}"
        --judge_gpu_memory "${JUDGE_GPU_MEMORY}"
        --judge_num_gpus "${JUDGE_NUM_GPUS}"
        "${ORIG_ARGS[@]}"
    )

    if [ "${INFERENCE_ONLY}" = "true" ]; then
        CMD+=(--inference_only)
        echo ""
        echo ">>> Stage 1 only: LoRA inference"
    elif [ "${SKIP_JUDGE}" = "true" ]; then
        CMD+=(--skip_judge)
        echo ""
        echo ">>> Stage 1 + regex detection (no LLM judge)"
    else
        echo ""
        echo ">>> Full pipeline: Stage 1 (LoRA inference) + Stage 1.5 (LLM judge)"
    fi

    "${CMD[@]}"
fi

echo ""
echo "=========================================="
echo "Done!"
echo "  Output: ${JUDGE_EXISTING:-${OUTPUT}}"
echo "=========================================="