#!/usr/bin/env bash
# Data-parallel eval: launch 8 single-GPU vLLM workers in parallel, merge shard outputs.
#
# Usage:
#   scripts/eval_hallu_dp.sh <MODEL_PATH_OR_NAME> <TAG> [bench=both|videohallucer|eventhallusion]
#
# Outputs:
#   outputs/eval/${TAG}_<bench>.jsonl   (merged predictions)
#   outputs/eval/${TAG}_<bench>.json    (merged summary metrics)
set -euo pipefail

ROOT=/mnt/local-fast/opd_zt
cd "$ROOT"
source .env
source .venv/bin/activate

MODEL=${1:?MODEL path or HF id required}
TAG=${2:?TAG required (e.g. opd_step468)}
BENCH_ARG=${3:-both}

OUT=outputs/eval
mkdir -p "$OUT" logs/eval

NUM_SHARDS=8
TP=1
MAX_LEN=12288
GPU_MEM=0.85

if [[ "$BENCH_ARG" == "both" ]]; then
  BENCHES=(videohallucer eventhallusion)
else
  BENCHES=("$BENCH_ARG")
fi

for BENCH in "${BENCHES[@]}"; do
  echo "=== $TAG $BENCH (DP=$NUM_SHARDS, TP=$TP) ==="
  STAMP=$(date +%H%M%S)
  PIDS=()
  for ((s=0; s<NUM_SHARDS; s++)); do
    LOG="logs/eval/${TAG}_${BENCH}_shard${s}_${STAMP}.log"
    CUDA_VISIBLE_DEVICES=$s python scripts/eval_hallu.py \
        --model "$MODEL" \
        --bench "$BENCH" \
        --out "$OUT/${TAG}_${BENCH}_shard${s}.jsonl" \
        --summary "$OUT/${TAG}_${BENCH}_shard${s}.json" \
        --tp $TP --max_model_len $MAX_LEN --gpu_mem $GPU_MEM \
        --shard $s --num_shards $NUM_SHARDS \
        > "$LOG" 2>&1 &
    PIDS+=("$!")
  done
  echo "Launched ${#PIDS[@]} shard workers: ${PIDS[*]}"

  FAIL=0
  for pid in "${PIDS[@]}"; do
    if ! wait "$pid"; then
      echo "[ERR] shard pid=$pid exited non-zero"
      FAIL=1
    fi
  done
  if (( FAIL )); then
    echo "[ERR] at least one shard failed; inspect logs/eval/${TAG}_${BENCH}_shard*_${STAMP}.log"
    exit 1
  fi

  # Merge shard predictions + re-score on the union.
  python scripts/merge_shards.py \
      --bench "$BENCH" \
      --in_glob "$OUT/${TAG}_${BENCH}_shard*.jsonl" \
      --out_jsonl "$OUT/${TAG}_${BENCH}.jsonl" \
      --out_summary "$OUT/${TAG}_${BENCH}.json"
  rm -f "$OUT/${TAG}_${BENCH}_shard"*.jsonl "$OUT/${TAG}_${BENCH}_shard"*.json
  echo "=== DONE $TAG $BENCH -> $OUT/${TAG}_${BENCH}.json ==="
done