#!/usr/bin/env bash
# =============================================================================
# Build Caption Targets: LLaVA Inference + LLM Judge + LLM Cleaning
# =============================================================================
# Produces a reusable caption_targets.json that any edit method can consume.
#
# Three stages:
#   Stage 1:   Run original LLaVA on all images -> raw captions
#   Stage 1.5: Regex coarse filter + LLM judge to confirm object mentions
#   Stage 2:   Use LLM to rewrite hallucinating captions (object removed)
#
# Data is loaded from HuggingFace by default (dataset auto-resolved from RELATION).
# Set CSV_PATH + IMAGE_DIR to use local files instead.
#
# Usage:
#   # Full pipeline for a specific relation
#   RELATION=kitchen_microwave bash experiment/scripts/data/run_build_caption_targets.sh
#
#   # Inference only (judge/clean later)
#   RELATION=kitchen_microwave INFERENCE_ONLY=true \
#       bash experiment/scripts/data/run_build_caption_targets.sh
#
#   # Judge an existing file (stage 1.5 only)
#   JUDGE_EXISTING=experiment/data/caption_targets_kitchen_microwave.json \
#       bash experiment/scripts/data/run_build_caption_targets.sh
#
#   # Clean an existing file (stage 2 only)
#   CLEAN_EXISTING=experiment/data/caption_targets_kitchen_microwave.json \
#       bash experiment/scripts/data/run_build_caption_targets.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
# =============================================================================
RELATION="${RELATION:-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:-}"

OUTPUT="${OUTPUT:-experiment/data/caption_targets_${RELATION}.json}"
MODEL="${MODEL:-llava-hf/llava-1.5-7b-hf}"
JUDGE_MODEL="${JUDGE_MODEL:-Qwen/Qwen3-8B}"
CLEANER_MODEL="${CLEANER_MODEL:-Qwen/Qwen3-8B}"
DEVICE="${DEVICE:-cuda}"
INFERENCE_ONLY="${INFERENCE_ONLY:-false}"
SKIP_JUDGE="${SKIP_JUDGE:-false}"
JUDGE_EXISTING="${JUDGE_EXISTING:-}"
CLEAN_EXISTING="${CLEAN_EXISTING:-}"
BATCH_SIZE="${BATCH_SIZE:-64}"
GPU_MEMORY="${GPU_MEMORY:-0.8}"
NUM_GPUS="${NUM_GPUS:-1}"
JUDGE_BATCH_SIZE="${JUDGE_BATCH_SIZE:-64}"
JUDGE_GPU_MEMORY="${JUDGE_GPU_MEMORY:-0.8}"
JUDGE_TP="${JUDGE_TP:-1}"
CLEANER_BATCH_SIZE="${CLEANER_BATCH_SIZE:-64}"
CLEANER_GPU_MEMORY="${CLEANER_GPU_MEMORY:-0.8}"
CLEANER_TP="${CLEANER_TP:-1}"

echo "=========================================="
echo "Build Caption Targets"
echo "=========================================="
echo "  Relation:       ${RELATION}"
echo "Data:"
if [ -n "${CSV_PATH}" ] && [ -n "${IMAGE_DIR}" ]; then
    echo "  CSV:            ${CSV_PATH}"
    echo "  Image dir:      ${IMAGE_DIR}"
else
    echo "  Dataset:        (auto from relation config)"
fi
echo "  GPUs:           ${CUDA_VISIBLE_DEVICES:-all}"
echo "  Output:         ${OUTPUT}"
echo "  LLaVA model:    ${MODEL} (transformers)"
echo "  Judge model:    ${JUDGE_MODEL} (vLLM)"
echo "  Cleaner model:  ${CLEANER_MODEL} (vLLM)"
echo "  Device:         ${DEVICE}"
echo "  Inference only: ${INFERENCE_ONLY}"
echo "  Skip judge:     ${SKIP_JUDGE}"
echo "  Judge existing: ${JUDGE_EXISTING:-none}"
echo "  Clean existing: ${CLEAN_EXISTING:-none}"
echo "=========================================="

if [ -n "${JUDGE_EXISTING}" ]; then
    # Stage 1.5 only: judge an existing file
    echo ""
    echo ">>> Stage 1.5: LLM judge of ${JUDGE_EXISTING}"
    python -m experiment.data.build_caption_targets \
        --relation "${RELATION}" \
        --judge_only "${JUDGE_EXISTING}" \
        --judge_model "${JUDGE_MODEL}" \
        --judge_batch_size "${JUDGE_BATCH_SIZE}" \
        --judge_gpu_memory "${JUDGE_GPU_MEMORY}" \
        --judge_tp "${JUDGE_TP}"

elif [ -n "${CLEAN_EXISTING}" ]; then
    # Stage 2 only: clean an existing file
    echo ""
    echo ">>> Stage 2: LLM cleaning of ${CLEAN_EXISTING}"
    python -m experiment.data.build_caption_targets \
        --relation "${RELATION}" \
        --clean "${CLEAN_EXISTING}" \
        --cleaner_model "${CLEANER_MODEL}" \
        --cleaner_batch_size "${CLEANER_BATCH_SIZE}" \
        --cleaner_gpu_memory "${CLEANER_GPU_MEMORY}" \
        --cleaner_tp "${CLEANER_TP}"
else
    # Build data args
    DATA_ARGS=(--relation "${RELATION}")
    if [ -n "${CSV_PATH}" ] && [ -n "${IMAGE_DIR}" ]; then
        DATA_ARGS+=(--csv "${CSV_PATH}" --image_dir "${IMAGE_DIR}")
    fi

    # Full pipeline or inference-only
    CMD=(
        python -m experiment.data.build_caption_targets
        "${DATA_ARGS[@]}"
        --output "${OUTPUT}"
        --model "${MODEL}"
        --judge_model "${JUDGE_MODEL}"
        --cleaner_model "${CLEANER_MODEL}"
        --device "${DEVICE}"
        --batch_size "${BATCH_SIZE}"
        --gpu_memory "${GPU_MEMORY}"
        --num_gpus "${NUM_GPUS}"
        --judge_batch_size "${JUDGE_BATCH_SIZE}"
        --judge_gpu_memory "${JUDGE_GPU_MEMORY}"
        --judge_tp "${JUDGE_TP}"
        --cleaner_batch_size "${CLEANER_BATCH_SIZE}"
        --cleaner_gpu_memory "${CLEANER_GPU_MEMORY}"
        --cleaner_tp "${CLEANER_TP}"
    )

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

    "${CMD[@]}"
fi

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