#!/usr/bin/env bash
# RAGTruth baselines, split by memory profile and run sequentially.
#
# The 12-detector single pass was killed by the Linux OOM killer at 47.9 GB
# resident. Cause: normalize_attentions() stacks all 32 layers into one dense
# (32, 32, N, N) float32 array on the HOST, and four consumers then build on it.
# That was written for v5.2's 44-token sequences; RAGTruth rows reach N ~ 1500+,
# where a single copy is 9+ GB and several coexist.
#
# So the work is split by whether a detector needs output_attentions at all:
#
#   PASS A  9 detectors, hidden states + logits only. Never allocates the big
#           tensor. This is the safe pass and it runs FIRST, so an OOM later
#           cannot cost the whole night.
#   PASS B  lookback alone. Needs attention, but is the single most important
#           attention baseline (our closest competitor), so it gets its own pass
#           with nothing else competing for memory.
#   PASS C  llm_check_attention, then icr_probe -- separately, because ICR
#           stacks the attention tensor a second time ON THE GPU and
#           mask_attention_blocks copies it again.
#
# Each pass writes its own summary and cache. A pass that dies does not stop the
# rest: `set -e` is deliberately NOT used.
#
#   bash scripts/run_ragtruth_baselines.sh

set -u -o pipefail

REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$REPO"

PY="${PY:-$REPO/.venv/bin/python}"
DATA="${DATA:-$HOME/datasets/hallucinations_custom/ragtruth_llama2_7b_chat.jsonl}"
SPLIT="${SPLIT:-$HOME/datasets/hallucinations_custom/ragtruth_llama2_7b_chat_split.json}"
MODEL="${MODEL:-NousResearch/Llama-2-7b-chat-hf}"
LAYER="${LAYER:-16}"
SEED="${SEED:-42}"
ROOT="${ROOT:-$REPO/outputs/rt_baselines}"

mkdir -p "$ROOT"
STATUS="$ROOT/status.tsv"; : > "$STATUS"

echo "### RAGTruth baselines (Llama-2-7b-chat, layer $LAYER, seed $SEED)"
echo "started: $(date -Is)"
echo "LatticeSpike marks to beat: AP 0.4103 / AUROC 0.8389 (random-AP floor 0.0625)"
echo

run_pass() {
  local name="$1"; shift
  local out="$ROOT/$name"
  mkdir -p "$out"
  echo "=== [$name] start $(date -Is)"
  local start rc elapsed
  start=$(date +%s)
  "$PY" -u scripts/train_faithful_baselines_dgx.py \
      --model "$MODEL" --allow-unpinned-model \
      --input "$DATA" --split-summary "$SPLIT" \
      --suite causal_matched \
      "$@" \
      --hidden-layer "$LAYER" --llm-check-attention-layer "$LAYER" \
      --seed "$SEED" --device cuda --dtype bfloat16 \
      --prompt-template raw --allow-unverified-replay --strict-labels \
      --feature-cache-dir "$ROOT/cache_$name" \
      --output "$out/predictions.jsonl" \
      --summary "$out/summary.json" \
      --save-dir "$out/models" \
      > "$out/run.log" 2>&1
  rc=$?; elapsed=$(( $(date +%s) - start ))
  if [[ $rc -eq 0 ]]; then
    echo "=== [$name] OK in ${elapsed}s"
    printf '%s\tOK\t%s\n' "$name" "$elapsed" >> "$STATUS"
  else
    # 137 = SIGKILL, which is what the OOM killer sends.
    local why="rc=$rc"
    [[ $rc -eq 137 ]] && why="OOM-KILLED"
    echo "=== [$name] FAILED ($why) after ${elapsed}s"
    tail -n 6 "$out/run.log" | sed "s/^/    [$name] /"
    printf '%s\tFAILED(%s)\t%s\n' "$name" "$why" "$elapsed" >> "$STATUS"
  fi
  echo
}

# ---- PASS A: no attention. Safe, runs first. -------------------------------
run_pass "a_no_attention" \
  --detector epr --detector wepr \
  --detector token_logprob \
  --detector llm_check_logit_entropy --detector llm_check_window_entropy \
  --detector llm_check_hidden \
  --detector saplma --detector hami --detector hidden_state_probe

# ---- PASS B: lookback alone. -----------------------------------------------
run_pass "b_lookback" --detector lookback

# ---- PASS C: the two heaviest, separately. ---------------------------------
run_pass "c_llm_check_attention" --detector llm_check_attention
run_pass "d_icr_probe" --detector icr_probe

echo "=== finished $(date -Is) ==="
sort "$STATUS"
echo
echo "Collecting whatever completed:"
"$PY" - "$ROOT" <<'PYEOF'
import json, sys, glob, os
root = sys.argv[1]
rows = []
for f in sorted(glob.glob(os.path.join(root, "*", "summary.json"))):
    try:
        s = json.load(open(f))
    except Exception:
        continue
    for r in s.get("eval", []):
        rows.append((r.get("auprc") or 0.0, r["method"], r.get("auroc") or 0.0,
                     r.get("n_metric_tokens")))
if not rows:
    print("  (no summaries yet)")
else:
    print(f"  {'method':<28}{'TEST AP':>9}{'AUROC':>9}{'tokens':>10}")
    for ap, m, au, n in sorted(rows, reverse=True):
        print(f"  {m:<28}{ap:>9.4f}{au:>9.4f}{n if n else '':>10}")
    print("\n  LatticeSpike: AP 0.4103  AUROC 0.8389  tokens 87678")
    print("  random-AP floor: 0.0625")
PYEOF