#!/usr/bin/env bash
# =============================================================================
# Magic — Zindi MSRH E2E reproduction driver
# =============================================================================
# One-command runner that regenerates `go.csv` from the 19 shipped LoRA adapters.
#
#   Usage:  bash scripts/run_all.sh
#
# Prereqs (see MODEL_CARD.md / environment.md):
#   * 8 × 80 GB GPUs (H100 / A100), CUDA driver ≥ 12.0
#   * conda envs `vllm` (inference) and `ensemble` (CPU)
#   * Base Qwen model snapshots downloaded to $ROOT/hub/{Qwen3.5-27B,Qwen3.6-27B,Qwen3-32B}
#   * SampleSubmission.csv + Test.csv in $ROOT/data/
#
# Outputs:
#   * $ROOT/predict_out/<cand>/generated_predictions.jsonl   (per adapter)
#   * $ROOT/candidate_csvs/<cand>.csv                        (per adapter)
#   * $ROOT/submission.csv                                   (final ensemble)
# =============================================================================
set -euo pipefail

ROOT="${ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}"
CKPT_DIR="${CKPT_DIR:-$ROOT/checkpoints}"
CAND_DIR="${CAND_DIR:-$ROOT/candidate_csvs}"
PRED_DIR="${PRED_DIR:-$ROOT/predict_out}"

VLLM_ENV="${VLLM_ENV:-vllm}"          # conda env with vllm + peft
ENSEMBLE_ENV="${ENSEMBLE_ENV:-ensemble}"  # conda env with rouge-score

mkdir -p "$CAND_DIR" "$PRED_DIR"

log()  { echo "[$(date +%H:%M:%S)]  $*"; }
step() { echo; echo "======  $*  ======"; }

# ---------------------------------------------------------------------------
step "STEP 1 / 3 : run vLLM inference for all 19 adapters"
# ---------------------------------------------------------------------------
log "using launcher $ROOT/scripts/launch_all_predicts.sh"
bash "$ROOT/scripts/launch_all_predicts.sh"

# ---------------------------------------------------------------------------
step "STEP 2 / 3 : convert per-adapter JSONL -> CSV"
# ---------------------------------------------------------------------------
CONVERTED=0; SKIPPED=0
for d in "$CKPT_DIR"/*/; do
  name=$(basename "$d")
  jsonl="$PRED_DIR/$name/generated_predictions.jsonl"
  out_csv="$CAND_DIR/$name.csv"
  if [ ! -s "$jsonl" ]; then
    log "  ! skip $name (jsonl missing)"
    SKIPPED=$((SKIPPED + 1)); continue
  fi
  python3 "$ROOT/scripts/jsonl_rowidx_to_csv.py" \
      --jsonl   "$jsonl" \
      --out_csv "$out_csv"
  CONVERTED=$((CONVERTED + 1))
done
log "converted=$CONVERTED  skipped=$SKIPPED"

if [ "$CONVERTED" -ne 19 ]; then
  echo "!! Expected 19 candidate CSVs, got $CONVERTED. Aborting before ensemble." >&2
  exit 2
fi

# ---------------------------------------------------------------------------
step "STEP 3 / 3 : build V2 medoid_ngram ensemble + MD5 verify vs go.csv"
# ---------------------------------------------------------------------------
python3 "$ROOT/scripts/build_ensemble.py"

# ---------------------------------------------------------------------------
step "DONE"
# ---------------------------------------------------------------------------
log "candidate CSVs :  $CAND_DIR"
log "final CSV      :  $ROOT/submission.csv"
log "shipped CSV    :  $ROOT/go.csv"
log ""
log "Compare (any diff should be minor float / vLLM non-determinism):"
log "  md5sum $ROOT/submission.csv $ROOT/go.csv"