#!/usr/bin/env bash
# engine2 Roformer VRAM spike (task #8) — does a SOTA Mel-/BS-Band Roformer fit
# on the 6 GB RTX 2060? Builds an isolated audio-separator venv, runs ONE pass on
# the Loituma source, and samples nvidia-smi throughout to capture PEAK VRAM.
#
# Self-contained + idempotent: re-running reuses the venv. All output → ./ (this
# dir), so results survive the session. Read SPIKE.log + result.json after.
set -uo pipefail

HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# Python 3.14 breaks audio-separator: beartype rejects the Roformer's
# `Callable | None` hints (run 1, 2026-06-23). Pin an older interpreter.
PYBIN="${SPIKE_PY:-python3.12}"
VENV="$HOME/.virtualenvs/audio-separator"
SRC="$HOME/Downloads/foundry/loituma_ievan_polkka_hqthspskzv8/source.webm"
OUT="$HERE/out"
LOG="$HERE/SPIKE.log"
VRAM_SAMPLES="$HERE/vram_samples.txt"
RESULT="$HERE/result.json"

mkdir -p "$OUT"
: > "$LOG"; : > "$VRAM_SAMPLES"
say(){ echo "[$(date +%H:%M:%S)] $*" | tee -a "$LOG"; }

say "=== engine2 Roformer VRAM spike ==="
say "GPU: $(nvidia-smi --query-gpu=name,memory.total --format=csv,noheader 2>/dev/null)"
say "source: $SRC ($(du -h "$SRC" | cut -f1))"

# ── 1. venv + audio-separator (GPU) ────────────────────────────────────────
# Rebuild unless the venv exists AND runs the pinned interpreter version.
WANT_VER="$("$PYBIN" -c 'import sys;print("%d.%d"%sys.version_info[:2])' 2>/dev/null)"
HAVE_VER="$("$VENV/bin/python" -c 'import sys;print("%d.%d"%sys.version_info[:2])' 2>/dev/null || true)"
if [ ! -x "$VENV/bin/audio-separator" ] || [ "$WANT_VER" != "$HAVE_VER" ]; then
  say "creating venv $VENV with $PYBIN ($WANT_VER) — was '${HAVE_VER:-none}'"
  rm -rf "$VENV"
  "$PYBIN" -m venv "$VENV" 2>&1 | tee -a "$LOG"
  say "pip install audio-separator[gpu] (this is the slow part)…"
  "$VENV/bin/pip" install --upgrade pip 2>&1 | tail -2 | tee -a "$LOG"
  "$VENV/bin/pip" install "audio-separator[gpu]" 2>&1 | tail -25 | tee -a "$LOG"
else
  say "venv already present, reusing"
fi

SEP="$VENV/bin/audio-separator"
if [ ! -x "$SEP" ]; then say "FATAL: audio-separator not installed"; exit 1; fi
say "version: $($SEP --version 2>&1 | head -1)"

# ── 2. background VRAM sampler (every 1s) ───────────────────────────────────
( while true; do
    nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null >> "$VRAM_SAMPLES"
    sleep 1
  done ) &
SAMPLER=$!
trap 'kill $SAMPLER 2>/dev/null' EXIT

# ── 3. run ONE Mel-Band Roformer pass, timed ───────────────────────────────
# Mel-Band Roformer (Kim) — strong 4-stem-ish vocal/inst SOTA; the model the
# README earmarked. audio-separator downloads weights on first use.
MODEL="model_bs_roformer_ep_317_sdr_12.9755.ckpt"
say "model: $MODEL  (downloads on first run)"
START=$(date +%s)
set +e
"$SEP" "$SRC" \
  --model_filename "$MODEL" \
  --output_dir "$OUT" \
  --output_format WAV \
  --log_level INFO 2>&1 | tee -a "$LOG"
RC=$?
set -e 2>/dev/null
END=$(date +%s)
ELAPSED=$((END - START))

kill $SAMPLER 2>/dev/null

# ── 4. summarize ────────────────────────────────────────────────────────────
PEAK=$(sort -n "$VRAM_SAMPLES" | tail -1)
TOTAL=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null | head -1)
STEMS=$(ls -1 "$OUT" 2>/dev/null | wc -l)
say "---"
say "exit code:        $RC"
say "elapsed:          ${ELAPSED}s (source 164s)"
say "peak VRAM:        ${PEAK:-?} / ${TOTAL:-?} MiB"
say "output files:     $STEMS"
ls -la "$OUT" 2>/dev/null | tee -a "$LOG"

# verdict: fits if it produced output AND peak < total (no OOM crash)
FITS="unknown"
if [ "$RC" = "0" ] && [ "${STEMS:-0}" -gt 0 ]; then FITS="yes"; else FITS="no"; fi

python3 - "$RC" "$ELAPSED" "${PEAK:-0}" "${TOTAL:-6144}" "$STEMS" "$FITS" "$MODEL" <<'PY' > "$RESULT"
import json, sys
rc, elapsed, peak, total, stems, fits, model = sys.argv[1:8]
print(json.dumps({
    "task": "#8 engine2 roformer VRAM spike",
    "gpu": "RTX 2060 Max-Q 6GB",
    "model": model,
    "exit_code": int(rc),
    "elapsed_s": int(elapsed),
    "source_s": 164,
    "peak_vram_mib": int(peak),
    "total_vram_mib": int(total),
    "vram_headroom_mib": int(total) - int(peak),
    "output_stems": int(stems),
    "fits_6gb": fits,
}, indent=2))
PY
say "=== verdict: fits_6gb=$FITS — see result.json ==="
cat "$RESULT" | tee -a "$LOG"
