#!/usr/bin/env bash
# Priority items 1, 3 and 7, in order, unattended.
#
#   1. LatticeSpike on v5.2, seeds 45-46  -> completes a five-seed mean there
#   3. SAPLMA / HaMI / ICR at seeds 42-46 on BOTH corpora
#   7. RAGTruth cross-task transfer grid (ours / probe / Lookback), batch 12
#
# Item 2 (RAGTruth seeds) is already done -- ablations/rt_bs12/ -- so item 1
# takes its place as the remaining seed work on our own model.
#
# ORDER IS DELIBERATE. Item 3 is refits over caches that already exist and takes
# minutes, so it could go first; it does not, because item 1 is the longer job
# and starting it first means the GPU is never idle waiting on a decision. Item
# 7 runs last because it is the one whose result is most likely to change what
# we do next, and it should not block the cheap gap-closing work behind it.
#
# WHY THESE THREE:
#   1 -- v5.2 currently reports three seeds against baselines that are
#        deterministic single runs. Five makes the error bar comparable to
#        RAGTruth's and removes "only three seeds" as an objection.
#   3 -- SAPLMA, HaMI and ICR train neural nets from random initialisation, so
#        their single-seed numbers have unmeasured spread. Every other baseline
#        in the table is deterministic (lbfgs or an untrained statistic) and
#        genuinely needs one run. These three do not.
#   7 -- the only remaining experiment that can produce a NEW finding rather
#        than a tighter version of an existing one. Lookback Lens argued
#        hidden-state probes overfit their training distribution; this tests
#        that on our corpus, where the probe is our closest competitor.
#
# All three are seed-only or refit work -- no feature extraction, nothing
# re-derived from the LLM. The signal and feature caches are keyed on model,
# prompt and layer, never on the seed or the split.
#
#   bash scripts/run_remaining.sh
#   ONLY=3 bash scripts/run_remaining.sh      # a single item
#   SKIP=1 bash scripts/run_remaining.sh      # skip one, run the rest

set -u -o pipefail
REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"; cd "$REPO"
PY="${PY:-$REPO/.venv/bin/python}"
ONLY="${ONLY:-}"
ITEMS="${ITEMS:-3 4 7}"
SEEDS="${SEEDS:-42 43 44 45 46}"

V52_DATA="$HOME/datasets/hallucinations_custom/hallu_dataset_v5.2_uncertain_eval.jsonl"
V52_SPLIT="$HOME/datasets/hallucinations_custom/hallu_v5_2_split.json"
V52_MODEL="meta-llama/Llama-3.2-1B-Instruct"
V52_REV="9213176726f574b556790deb65791e0c5aa438b6"
RT_DATA="$HOME/datasets/hallucinations_custom/ragtruth_llama2_7b_chat.jsonl"
RT_SPLIT="$HOME/datasets/hallucinations_custom/ragtruth_llama2_7b_chat_split.json"
RT_MODEL="NousResearch/Llama-2-7b-chat-hf"

ROOT="${ROOT:-$REPO/outputs/remaining}"
mkdir -p "$ROOT"; STATUS="$ROOT/status.tsv"; : > "$STATUS"

want () {
  if [[ -n "$ONLY" ]]; then [[ "$ONLY" == "$1" ]]; return; fi
  [[ " $ITEMS " == *" $1 "* ]]
}
say  () { echo; echo "############ $* ############"; echo; }
rec  () { printf '%s\t%s\t%s\n' "$1" "$2" "$3" >> "$STATUS"; }

# ---------------------------------------------------------------- item 1 ----
# v5.2 config is BEST_MODEL.md's: batch 4, tau 0.25, lam_rank 12, aux losses
# zero. NOT batch 12 -- that was established on RAGTruth and the batch axis has
# never been varied on v5.2 (all 71 runs across r0-r6 used batch 4). Changing it
# here would silently alter the corpus where we currently sweep every metric.
if want 1; then
  say "ITEM 1  LatticeSpike v5.2, seeds 45-46 (batch 4)"
  pids=()
  for SEED in 45 46; do
    OUT="$ROOT/v52_s${SEED}"
    if [[ -f "$OUT/summary.json" ]]; then echo "[skip] v52 s$SEED"; rec "v52_s$SEED" SKIP 0; continue; fi
    mkdir -p "$OUT"; echo "=== v52 seed $SEED  $(date +%H:%M:%S)"
    start=$(date +%s)
    "$PY" -u scripts/train_latticespike_llm_eval.py \
        --input "$V52_DATA" --split-summary "$V52_SPLIT" \
        --model "$V52_MODEL" --model-revision "$V52_REV" \
        --prompt-template tanazzah-llama-chat \
        --allow-unverified-replay --skip-extract --tie-box-token-emb \
        --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 "$SEED" \
        --signals-dir outputs/latticespike_v5_2_signals \
        --out-dir "$OUT" --summary "$OUT/summary.json" > "$OUT/run.log" 2>&1 &
    pids+=($!)
  done
  wait || true
  for SEED in 45 46; do
    OUT="$ROOT/v52_s${SEED}"
    [[ -f "$OUT/summary.json" ]] && rec "v52_s$SEED" OK 0 || { rec "v52_s$SEED" FAILED 0; tail -4 "$OUT/run.log" 2>/dev/null | sed 's/^/    /'; }
  done
  "$PY" - <<'PYEOF' || true
import glob, json, os, statistics
paths = {"s42": "outputs/sweep_r3_5/03_tau_0p25/summary.json",
         "s43": "outputs/sweep_r4/01_tau025_s43/summary.json",
         "s44": "outputs/sweep_r4/02_tau025_s44/summary.json"}
for f in sorted(glob.glob("outputs/remaining/v52_s*/summary.json")):
    paths[os.path.basename(os.path.dirname(f)).replace("v52_", "")] = f
ap = {k: (json.load(open(p)).get("eval") or {}).get("auprc")
      for k, p in paths.items() if os.path.exists(p)}
ap = {k: v for k, v in ap.items() if v is not None}
if len(ap) > 1:
    print("\n  v5.2 AP over %d seeds: %s" % (len(ap), {k: round(v,4) for k,v in sorted(ap.items())}))
    print("  mean %.4f  sd %.4f   (lookback 0.2305, probe 0.2286)"
          % (statistics.mean(ap.values()), statistics.stdev(ap.values())))
PYEOF
fi

# ---------------------------------------------------------------- item 3 ----
# Refits only: the feature caches for these detectors already exist from the
# original passes, and the cache key covers model/prompt/layer/detector -- not
# the seed. ICR gets its own invocation because it stacks the attention tensor a
# second time on the GPU, which is why the first RAGTruth run split it out.
if want 3; then
  say "ITEM 3  SAPLMA / HaMI / ICR at seeds $SEEDS, both corpora"
  bl_run () {
    local name="$1" seed="$2"; shift 2
    local out="$ROOT/$name"
    if [[ -f "$out/summary.json" ]]; then echo "[skip] $name"; rec "$name" SKIP 0; return; fi
    mkdir -p "$out"; local start; start=$(date +%s)
    "$PY" -u scripts/train_faithful_baselines_dgx.py "$@" \
        --suite causal_matched --seed "$seed" \
        --device cuda --dtype bfloat16 \
        --allow-unverified-replay --strict-labels \
        --output "$out/predictions.jsonl" --summary "$out/summary.json" \
        --save-dir "$out/models" > "$out/run.log" 2>&1
    local rc=$? el=$(( $(date +%s) - start ))
    if [[ $rc -eq 0 ]]; then rec "$name" OK "$el"; echo "  [$name] ok ${el}s"
    else rec "$name" "FAILED(rc=$rc)" "$el"; echo "  [$name] FAILED"; tail -4 "$out/run.log" | sed 's/^/    /'; fi
  }
  for SEED in $SEEDS; do
    echo "--- seed $SEED"
    bl_run "v52_sapl_hami_s$SEED" "$SEED" \
      --model "$V52_MODEL" --model-revision "$V52_REV" \
      --input "$V52_DATA" --split-summary "$V52_SPLIT" \
      --prompt-template tanazzah-llama-chat --hidden-layer 8 --llm-check-attention-layer 8 \
      --detector saplma --detector hami --feature-cache-dir "$ROOT/cache_v52_cheap"
    bl_run "v52_icr_s$SEED" "$SEED" \
      --model "$V52_MODEL" --model-revision "$V52_REV" \
      --input "$V52_DATA" --split-summary "$V52_SPLIT" \
      --prompt-template tanazzah-llama-chat --hidden-layer 8 --llm-check-attention-layer 8 \
      --detector icr_probe --feature-cache-dir "$ROOT/cache_v52_icr"
    bl_run "rt_sapl_hami_s$SEED" "$SEED" \
      --model "$RT_MODEL" --allow-unpinned-model \
      --input "$RT_DATA" --split-summary "$RT_SPLIT" \
      --prompt-template raw --hidden-layer 16 --llm-check-attention-layer 16 \
      --detector saplma --detector hami --feature-cache-dir "$ROOT/cache_rt_cheap"
    bl_run "rt_icr_s$SEED" "$SEED" \
      --model "$RT_MODEL" --allow-unpinned-model \
      --input "$RT_DATA" --split-summary "$RT_SPLIT" \
      --prompt-template raw --hidden-layer 16 --llm-check-attention-layer 16 \
      --detector icr_probe --feature-cache-dir "$ROOT/cache_rt_icr"
  done
  "$PY" - "$ROOT" <<'PYEOF' || true
import glob, json, os, statistics, sys, collections
root = sys.argv[1]
agg = collections.defaultdict(list)
for f in sorted(glob.glob(os.path.join(root, "*_s4*", "summary.json"))):
    tag = os.path.basename(os.path.dirname(f))
    corpus = "v5.2" if tag.startswith("v52") else "RAGTruth"
    for r in json.load(open(f)).get("eval") or []:
        if r.get("auprc") is not None:
            agg[(corpus, r["method"])].append(r["auprc"])
if agg:
    print("\n  %-10s%-22s%10s%10s%7s" % ("corpus", "method", "AP mean", "sd", "seeds"))
    print("  " + "-" * 59)
    for (c, m), v in sorted(agg.items()):
        sd = statistics.stdev(v) if len(v) > 1 else 0.0
        print("  %-10s%-22s%10.4f%10.4f%7d" % (c, m, statistics.mean(v), sd, len(v)))
    print("\n  Non-zero sd here is the point: these three train neural nets from")
    print("  random init, so a single-seed number for them was never safe.")
PYEOF
fi

# ---------------------------------------------------------------- item 4 ----
# Same features, same caches, same seeds -- only the training unit differs, so
# these are pure refits reusing item 3's cache directories. The faithful
# response-supervised runs are untouched; these land under separate method names
# (saplma_token_supervised, icr_probe_token_supervised) and are a CONTROL.
#
# The question: SAPLMA and ICR are published as RESPONSE-level classifiers, so
# we train them on response labels while our model gets token labels. Both score
# near the AP floor. Is that the method, or the supervision we gave them?
if want 4; then
  say "ITEM 4  token-supervised SAPLMA / ICR at seeds $SEEDS, both corpora"
  for SEED in $SEEDS; do
    echo "--- seed $SEED"
    bl_run "v52_tok_sapl_s$SEED" "$SEED" \
      --model "$V52_MODEL" --model-revision "$V52_REV" \
      --input "$V52_DATA" --split-summary "$V52_SPLIT" \
      --prompt-template tanazzah-llama-chat --hidden-layer 8 --llm-check-attention-layer 8 \
      --detector saplma --saplma-supervision token \
      --feature-cache-dir "$ROOT/cache_v52_cheap"
    bl_run "v52_tok_icr_s$SEED" "$SEED" \
      --model "$V52_MODEL" --model-revision "$V52_REV" \
      --input "$V52_DATA" --split-summary "$V52_SPLIT" \
      --prompt-template tanazzah-llama-chat --hidden-layer 8 --llm-check-attention-layer 8 \
      --detector icr_probe --icr-supervision token \
      --feature-cache-dir "$ROOT/cache_v52_icr"
    bl_run "rt_tok_sapl_s$SEED" "$SEED" \
      --model "$RT_MODEL" --allow-unpinned-model \
      --input "$RT_DATA" --split-summary "$RT_SPLIT" \
      --prompt-template raw --hidden-layer 16 --llm-check-attention-layer 16 \
      --detector saplma --saplma-supervision token \
      --feature-cache-dir "$ROOT/cache_rt_cheap"
    bl_run "rt_tok_icr_s$SEED" "$SEED" \
      --model "$RT_MODEL" --allow-unpinned-model \
      --input "$RT_DATA" --split-summary "$RT_SPLIT" \
      --prompt-template raw --hidden-layer 16 --llm-check-attention-layer 16 \
      --detector icr_probe --icr-supervision token \
      --feature-cache-dir "$ROOT/cache_rt_icr"
  done
  "$PY" scripts/collect_supervision_control.py "$ROOT" || true
fi

# ---------------------------------------------------------------- item 7 ----
if want 7; then
  say "ITEM 7  RAGTruth cross-task transfer grid (batch 12)"
  bash scripts/run_ragtruth_crosstask.sh && rec crosstask OK 0 || rec crosstask FAILED 0
fi

echo; echo "=== all done $(date -Is) ==="; sort "$STATUS"