#!/usr/bin/env bash
# Wait for the tune sweep to finish, then run the queued work unattended.
#
# The sweep holds the GPU for hours and the next two jobs both need it, so this
# blocks until it is genuinely done rather than racing it. "Done" means no
# LatticeSpike trainer is alive AND the tune session is gone -- either alone is
# not enough, because the session outlives the last arm briefly while it writes
# the summary table, and a trainer can be alive during a gap between waves.
#
# Steps run in order and a failure does not stop the rest: each is independently
# useful, and losing the TriviaQA pilot should not cost the seed-44 run that
# already succeeded.
#
#   tmux new -s queue 'bash scripts/queue_after_tune.sh 2>&1 | tee outputs/queue.log'
#
# Skip individual steps:
#   SKIP_SEED44=1 bash scripts/queue_after_tune.sh
#   SKIP_TRIVIAQA=1 bash scripts/queue_after_tune.sh

set -u -o pipefail
REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"; cd "$REPO"
PY="${PY:-$REPO/.venv/bin/python}"
POLL="${POLL:-120}"
WAIT_FOR="${WAIT_FOR:-tune}"
PILOT_N="${PILOT_N:-200}"

DATA="$HOME/datasets/hallucinations_custom/ragtruth_llama2_7b_chat.jsonl"
SPLIT="$HOME/datasets/hallucinations_custom/ragtruth_llama2_7b_chat_split.json"
SIGNALS="outputs/ragtruth_l2_7b_signals"

say () { printf '\n=== %s  %s\n' "$1" "$(date -Is)"; }

# --- wait -------------------------------------------------------------------
say "waiting for the '$WAIT_FOR' session and all LatticeSpike trainers to exit"
waited=0
while true; do
  session_alive=0; trainer_alive=0
  tmux has-session -t "$WAIT_FOR" 2>/dev/null && session_alive=1
  pgrep -f train_latticespike_llm_eval.py > /dev/null 2>&1 && trainer_alive=1
  if (( session_alive == 0 && trainer_alive == 0 )); then
    break
  fi
  # Report every ~30 min so a long wait does not look like a hang.
  if (( waited % 900 == 0 )); then
    running=$(pgrep -cf train_latticespike_llm_eval.py 2>/dev/null || echo 0)
    printf '    [%s] session=%s trainers=%s  (waited %dm)\n' \
      "$(date +%H:%M)" "$session_alive" "$running" "$((waited / 60))"
  fi
  sleep "$POLL"; waited=$(( waited + POLL ))
done
say "clear after $((waited / 60))m -- starting queued work"

# --- 1. the sweep's own ranked table ----------------------------------------
# The sweep prints this itself, but only if it exited cleanly; reprinting costs
# nothing and means the queue log is self-contained.
say "sweep results"
if [[ -d outputs/rt_tune ]]; then
  "$PY" - outputs/rt_tune <<'PYEOF' || true
import json, glob, os, sys
root = sys.argv[1]
rows = []
for f in sorted(glob.glob(os.path.join(root, "*", "summary.json"))):
    s = json.load(open(f))
    rows.append((
        (s.get("validation_eval") or {}).get("auprc", 0),
        (s.get("eval") or {}).get("auprc", 0),
        (s.get("eval") or {}).get("auroc", 0),
        (s.get("checkpoint_selection") or {}).get("best_epoch"),
        (s.get("training") or {}).get("epochs"),
        os.path.basename(os.path.dirname(f)),
    ))
rows.sort(reverse=True)
ctrl = next((r for r in rows if r[5] == "ctrl_s42"), None)
base = ctrl[0] if ctrl else None
print("%-12s%9s%9s%9s%9s%8s" % ("arm", "VAL AP", "dVAL", "TEST AP", "AUROC", "epoch"))
print("-" * 56)
for v, t, a, be, ep, n in rows:
    d = "--" if base is None else "%+.4f" % (v - base)
    print("%-12s%9.4f%9s%9.4f%9.4f%7s%s" % (n, v, d, t, a, be, "*" if be == ep else " "))
PYEOF
fi

# --- 2. seed 44 of the control ----------------------------------------------
# The two-seed spread is 0.0091 and our AP lead over the probe is 0.0074, so the
# third seed is what decides whether that lead is a result. Runs on the current
# config: if a sweep arm wins, the winner needs its own seeds regardless, and
# this still stands as the control's third point.
if [[ "${SKIP_SEED44:-0}" == "1" ]]; then
  say "seed 44 skipped"
elif [[ -f outputs/rt_ctrl_s44/summary.json ]]; then
  say "seed 44 already done"
else
  say "seed 44 (control config)"
  mkdir -p outputs/rt_ctrl_s44
  "$PY" -u scripts/train_latticespike_llm_eval.py \
      --input "$DATA" --split-summary "$SPLIT" \
      --model NousResearch/Llama-2-7b-chat-hf --allow-unpinned-model \
      --prompt-template raw --allow-unverified-replay --tie-box-token-emb \
      --skip-extract \
      --lam-contain 0 --lam-vol 0 --lam-recon 0 --lam-mrl 0 \
      --rank-objective ap --lam-rank 12 --ap-tau 0.25 \
      --epochs 10 --seed 44 \
      --signals-dir "$SIGNALS" \
      --out-dir outputs/rt_ctrl_s44 --summary outputs/rt_ctrl_s44/summary.json \
      > outputs/rt_ctrl_s44/run.log 2>&1 \
    && echo "    seed 44 OK" \
    || { echo "    seed 44 FAILED"; tail -4 outputs/rt_ctrl_s44/run.log | sed 's/^/    /'; }

  # Three seeds is the point of this, so print the spread as soon as it exists.
  "$PY" - <<'PYEOF' || true
import json, os
paths = {"s42": "outputs/rt_tune/ctrl_s42/summary.json",
         "s43": "outputs/rt_tune/ctrl_s43/summary.json",
         "s44": "outputs/rt_ctrl_s44/summary.json"}
ap = {}
for k, p in paths.items():
    if os.path.exists(p):
        ap[k] = (json.load(open(p)).get("eval") or {}).get("auprc")
if len(ap) >= 2:
    vals = [v for v in ap.values() if v is not None]
    mean = sum(vals) / len(vals)
    var = sum((v - mean) ** 2 for v in vals) / (len(vals) - 1) if len(vals) > 1 else 0
    print("\n    control TEST AP: " + "  ".join("%s=%.4f" % kv for kv in sorted(ap.items())))
    print("    mean %.4f  sd %.4f  over %d seeds" % (mean, var ** 0.5, len(vals)))
    print("    marks to beat: probe 0.4075, lookback 0.4036")
PYEOF
fi

# --- 3. TriviaQA pilot ------------------------------------------------------
# Deliberately a pilot. The 3-judge pass has never been timed here, and its
# output decides whether the corpus is usable at all -- if prevalence comes back
# near 0 or near 1 there is nothing to detect. Better to learn that from 200
# rows than from an overnight run on 3,000.
if [[ "${SKIP_TRIVIAQA:-0}" == "1" ]]; then
  say "TriviaQA pilot skipped"
else
  say "TriviaQA pilot (N=$PILOT_N)"
  MAX_ITEMS="$PILOT_N" bash scripts/run_triviaqa_pipeline.sh \
    && echo "    pilot OK" \
    || echo "    pilot FAILED -- likely missing judge models; see log above"
fi

say "queue finished"