#!/bin/bash
# =============================================================================
# Triadic Object-State Consistency (TOSC) Benchmark — runner
# =============================================================================
# Steps:
#   1. Run VLM captioning on origin / removed / replaced images (multi-GPU).
#   2. Merge the per-GPU answer shards.
#   3. Extract objects from captions and compute the TOSC metrics
#      (OPA / RCA / RUA / TOSC + fine-grained failure metrics).
#
# This repo is self-contained: inference (step 1) uses the vendored `llava`
# package at the repo root (`llava/eval/model_vqa.py`), and scoring (step 3)
# uses eval/eval_tosc.py. Nothing outside this repo is required at runtime.
#
# Prerequisites:
#   - python eval/build_index.py        (run once to create benchmark/*.jsonl)
#   - pip install -r requirements.txt   (torch / transformers / peft / ...;
#                                        flash-attn for the Qwen path)
# =============================================================================

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="${REPO_ROOT:-$(dirname "$SCRIPT_DIR")}"

# ==================== Model preset ====================
# Uncomment ONE block, or override MODEL_NAME / LORA_NAME / OUTPUT_NAME via env.
# model_vqa.py auto-routes: names containing "llava"+"1.5" use the LLaVA path,
# everything else uses the HuggingFace/Qwen path (which ignores --conv-mode).

# --- LLaVA-1.5-7B (TSA-DPO) ---
# MODEL_NAME="liuhaotian/llava-v1.5-7b"
# LORA_NAME="$REPO_ROOT/lora/LLaVA_v1_5_7b-TSA-DPO"
# OUTPUT_NAME="LLaVA_v1_5_7b-TSA-DPO"

# --- LLaVA-1.5-13B (TSA-DPO) ---
# MODEL_NAME="liuhaotian/llava-v1.5-13b"
# LORA_NAME="$REPO_ROOT/lora/LLaVA_v1_5_13b-TSA-DPO"
# OUTPUT_NAME="LLaVA_v1_5_13b-TSA-DPO"

# --- Qwen2-VL-7B (TSA-DPO) ---
# MODEL_NAME="Qwen/Qwen2-VL-7B-Instruct"
# LORA_NAME="$REPO_ROOT/lora/Qwen2_VL_7B-TSA-DPO"
# OUTPUT_NAME="Qwen2_VL_7B-TSA-DPO"

# --- Qwen2.5-VL-7B (TSA-DPO) [default] ---
MODEL_NAME="${MODEL_NAME:-Qwen/Qwen2.5-VL-7B-Instruct}"
LORA_NAME="${LORA_NAME:-$REPO_ROOT/lora/Qwen2_5_VL_7B-TSA-DPO}"
OUTPUT_NAME="${OUTPUT_NAME:-Qwen2_5_VL_7B-TSA-DPO}"

# To evaluate a base model without LoRA, set LORA_NAME="".

# ==================== Inference config ====================
GPU_LIST="${GPU_LIST:-0}"            # e.g. "0,1,2,3" for multi-GPU
CONV_MODE="${CONV_MODE:-vicuna_v1}"  # used by the LLaVA-1.5 path only
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-512}"
TEMPERATURE="${TEMPERATURE:-0}"      # 0 = greedy / deterministic
SEED="${SEED:-42}"

# ==================== Data / output config ====================
QUESTION_FILE="${QUESTION_FILE:-$REPO_ROOT/benchmark/TOSC_dataset.jsonl}"
INSERTIONS_FILE="${INSERTIONS_FILE:-$REPO_ROOT/benchmark/insertions.jsonl}"
IMAGE_FOLDER="${IMAGE_FOLDER:-$REPO_ROOT}"   # resolves the relative image_path in the jsonl
OUTPUT_DIR="${OUTPUT_DIR:-${TOSC_RESULTS:-/home/yilin/tmp/tosc_results}}"
ANSWERS_DIR="$OUTPUT_DIR/answers"

if [[ ! -f "$QUESTION_FILE" ]]; then
    echo "Question file not found: $QUESTION_FILE" >&2
    echo "Run 'python eval/build_index.py' first." >&2
    exit 1
fi

mkdir -p "$ANSWERS_DIR"

echo "=========================================="
echo "Step 1: VLM captioning"
echo "=========================================="
echo "Model:        $MODEL_NAME"
echo "LoRA:         ${LORA_NAME:-<none>}"
echo "Output name:  $OUTPUT_NAME"
echo "Question file:$QUESTION_FILE"
echo "Image folder: $IMAGE_FOLDER"
echo "GPU list:     $GPU_LIST"
echo "llava pkg:    $REPO_ROOT/llava (vendored)"

IFS=',' read -ra GPULIST <<<"$GPU_LIST"
CHUNKS=${#GPULIST[@]}

for IDX in $(seq 0 $((CHUNKS - 1))); do
    CHUNK_FILE="$ANSWERS_DIR/${OUTPUT_NAME}_${CHUNKS}_${IDX}.jsonl"
    rm -f "$CHUNK_FILE"
    echo "Launching GPU ${GPULIST[$IDX]} for chunk $IDX/$CHUNKS"
    CUDA_VISIBLE_DEVICES="${GPULIST[$IDX]}" \
    PYTHONPATH="$REPO_ROOT${PYTHONPATH:+:$PYTHONPATH}" \
        python -m llava.eval.model_vqa \
            --model-name "$MODEL_NAME" \
            --lora-name "$LORA_NAME" \
            --image-folder "$IMAGE_FOLDER" \
            --question-file "$QUESTION_FILE" \
            --answers-file "$CHUNK_FILE" \
            --temperature "$TEMPERATURE" \
            --num-chunks "$CHUNKS" \
            --chunk-idx "$IDX" \
            --max-new-tokens "$MAX_NEW_TOKENS" \
            --seed "$SEED" \
            --conv-mode "$CONV_MODE" &
    echo "Started GPU ${GPULIST[$IDX]} at PID $!"
    sleep 1
done

wait
echo "Captioning completed."

echo "=========================================="
echo "Step 2: Merging caption shards"
echo "=========================================="

FINAL_OUTPUT_FILE="$ANSWERS_DIR/${OUTPUT_NAME}.jsonl"
true >"$FINAL_OUTPUT_FILE"

for IDX in $(seq 0 $((CHUNKS - 1))); do
    CHUNK_FILE="$ANSWERS_DIR/${OUTPUT_NAME}_${CHUNKS}_${IDX}.jsonl"
    if [[ -f "$CHUNK_FILE" ]]; then
        echo "Concatenating $CHUNK_FILE"
        cat "$CHUNK_FILE" >>"$FINAL_OUTPUT_FILE"
        rm "$CHUNK_FILE"
    else
        echo "Missing chunk file: $CHUNK_FILE" >&2
        exit 1
    fi
done
echo "Merged captions saved to: $FINAL_OUTPUT_FILE"

echo "=========================================="
echo "Step 3: Evaluating TOSC"
echo "=========================================="

EVAL_OUTPUT_FILE="$OUTPUT_DIR/${OUTPUT_NAME}_eval.json"

python "$SCRIPT_DIR/eval_tosc.py" \
    --inference_file "$FINAL_OUTPUT_FILE" \
    --dataset_file "$QUESTION_FILE" \
    --insertions_file "$INSERTIONS_FILE" \
    --save_file "$EVAL_OUTPUT_FILE"

echo "=========================================="
echo "TOSC evaluation completed."
echo "Captions: $FINAL_OUTPUT_FILE"
echo "Results:  $EVAL_OUTPUT_FILE"
echo "Summary:  ${EVAL_OUTPUT_FILE%.json}_summary.json"
echo "=========================================="