#!/usr/bin/env bash
# One script: every method, every seed, one dataset, one comparison table.
#
#   bash scripts/run_final_eval.sh v52
#   bash scripts/run_final_eval.sh ragtruth
#   SEEDS="42 43 44 45 46" bash scripts/run_final_eval.sh v52
#
# Why this is one script and not one process: LatticeSpike reads a signal cache
# (--signals-dir) while the baselines read a feature cache (--feature-cache-dir),
# and the feature-cache key includes the detector set, so detector groups with
# different memory profiles need different caches. Those are separate
# invocations by construction. Sequencing them is a shell problem, and this is
# the shell.
#
# The seed is NOT part of any cache key -- extraction depends only on the model,
# prompt, layer and detector set. So the first seed pays for extraction and
# every later seed is a cache hit that repeats only the fits. Reseeding is
# minutes for baselines.
#
# IDEMPOTENT: any (method, seed) whose summary already exists is skipped, so
# re-running fills gaps rather than redoing work. Delete a summary to force it.

set -u -o pipefail

REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$REPO"
PY="${PY:-$REPO/.venv/bin/python}"
DATASET="${1:-}"
SEEDS="${SEEDS:-42 43 44}"

case "$DATASET" in
  v52)
    INPUT="$HOME/datasets/hallucinations_custom/hallu_dataset_v5.2_uncertain_eval.jsonl"
    SPLIT="$HOME/datasets/hallucinations_custom/hallu_v5_2_split.json"
    MODEL="meta-llama/Llama-3.2-1B-Instruct"
    MODEL_ARGS="--model-revision 9213176726f574b556790deb65791e0c5aa438b6"
    TEMPLATE="tanazzah-llama-chat"
    LAYER=8
    LS_SIGNALS="outputs/latticespike_v5_2_signals"
    ;;
  ragtruth)
    INPUT="$HOME/datasets/hallucinations_custom/ragtruth_llama2_7b_chat.jsonl"
    SPLIT="$HOME/datasets/hallucinations_custom/ragtruth_llama2_7b_chat_split.json"
    MODEL="NousResearch/Llama-2-7b-chat-hf"
    MODEL_ARGS="--allow-unpinned-model"
    TEMPLATE="raw"
    LAYER=16
    LS_SIGNALS="outputs/ragtruth_l2_7b_signals"
    ;;
  *)
    echo "usage: $0 {v52|ragtruth}   (env: SEEDS=\"42 43 44\")" >&2
    exit 2
    ;;
esac

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

# Detector groups. Split by memory profile, NOT preference: attention detectors
# each hold a per-layer float32 view of a (32, N, N) tensor, and on RAGTruth's
# longest rows running them together is what the OOM killer objected to.
GROUPS=(
  "cheap|--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"
  "lookback|--detector lookback"
  "attn|--detector llm_check_attention"
  "icr|--detector icr_probe"
)

echo "### FINAL EVAL — dataset=$DATASET  seeds=$SEEDS"
echo "started: $(date -Is)"
echo

record() { printf '%s\t%s\t%s\n' "$1" "$2" "$3" >> "$STATUS"; }

for SEED in $SEEDS; do
  # ---------- LatticeSpike ----------
  OUT="$ROOT/latticespike_s${SEED}"
  if [[ -f "$OUT/summary.json" ]]; then
    echo "[skip] latticespike seed $SEED (summary exists)"
    record "latticespike_s$SEED" SKIP 0
  else
    echo "=== latticespike seed $SEED  $(date -Is)"
    mkdir -p "$OUT"; start=$(date +%s)
    # shellcheck disable=SC2086
    "$PY" -u scripts/train_latticespike_llm_eval.py \
        --input "$INPUT" --split-summary "$SPLIT" \
        --model "$MODEL" $MODEL_ARGS --prompt-template "$TEMPLATE" \
        --allow-unverified-replay --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 "$LS_SIGNALS" \
        --out-dir "$OUT" --summary "$OUT/summary.json" \
        > "$OUT/run.log" 2>&1
    rc=$?; el=$(( $(date +%s) - start ))
    [[ $rc -eq 0 ]] && record "latticespike_s$SEED" OK "$el" \
                    || { record "latticespike_s$SEED" "FAILED(rc=$rc)" "$el"; tail -5 "$OUT/run.log" | sed 's/^/    /'; }
    # predictions/results land on hard-coded default paths; keep a copy that
    # the next run cannot clobber.
    cp -f outputs/latticespike_llama32_predictions.jsonl "$OUT/predictions.jsonl" 2>/dev/null || true
    cp -f outputs/latticespike_llama32_results.md "$OUT/results.md" 2>/dev/null || true
  fi

  # ---------- baselines, by group ----------
  for entry in "${GROUPS[@]}"; do
    name="${entry%%|*}"; flags="${entry#*|}"
    OUT="$ROOT/baselines_${name}_s${SEED}"
    if [[ -f "$OUT/summary.json" ]]; then
      echo "[skip] baselines/$name seed $SEED (summary exists)"
      record "baselines_${name}_s$SEED" SKIP 0
      continue
    fi
    echo "=== baselines/$name seed $SEED  $(date -Is)"
    mkdir -p "$OUT"; start=$(date +%s)
    # shellcheck disable=SC2086
    "$PY" -u scripts/train_faithful_baselines_dgx.py \
        --model "$MODEL" $MODEL_ARGS \
        --input "$INPUT" --split-summary "$SPLIT" \
        --suite causal_matched $flags \
        --hidden-layer "$LAYER" --llm-check-attention-layer "$LAYER" \
        --seed "$SEED" --device cuda --dtype bfloat16 \
        --prompt-template "$TEMPLATE" --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=$?; el=$(( $(date +%s) - start ))
    if [[ $rc -eq 0 ]]; then record "baselines_${name}_s$SEED" OK "$el"
    else
      why="rc=$rc"; [[ $rc -eq 137 ]] && why="OOM-KILLED"
      record "baselines_${name}_s$SEED" "FAILED($why)" "$el"
      tail -5 "$OUT/run.log" | sed 's/^/    /'
    fi
  done
done

echo
echo "=== finished $(date -Is) ==="
sort "$STATUS"
echo
"$PY" scripts/collect_final_eval.py "$ROOT" | tee "$ROOT/RESULTS.md"