#!/usr/bin/env bash
# Multi-layer hidden-state experiment (AWS A10G box).
#
# Question: the box encoder currently reads ONE hidden layer (8 of 16) and
# discards the other fifteen. Probing work (SAPLMA, INSIDE, semantic entropy
# probes) consistently finds layer choice matters and that different depths carry
# complementary truthfulness signal. Does reading several layers help?
#
# Method: read layers 4, 8, 12, 16 and learn an ELMo-style scalar mixture over
# them (softmax weights + scale). Deliberately NOT concatenation -- that would
# take the box encoder's input from 2048 to 8192 and confound any gain with a
# much larger input projection. This adds K+1 parameters and leaves the encoder
# input width unchanged, so the comparison isolates the extra information.
#
# Baseline for comparison is the 5090 result in ablations/BEST_MODEL.md:
#   VAL AP 0.2683 +/- 0.0088, TEST AP 0.2741 +/- 0.0238 (3 seeds, single layer 8)
#
# NOTE ON COST: this needs a fresh extraction (EXTRACTION_CODE_VERSION bumped to
# 5 and h_box gains a layer axis). The first run below extracts all 9,817 records
# -- roughly 6-7 hours on this box -- and seeds 43/44 then reuse that cache. The
# cache is written incrementally and validated per record, so an interrupted run
# resumes rather than restarting.
#
#   scripts/run_multilayer_aws.sh

set -u -o pipefail

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

PY="${PY:-$REPO/.venv/bin/python}"
DATA="${DATA:-$REPO/outputs/hallu_v5_2_uncertain_eval.jsonl}"
SPLIT="${SPLIT:-$REPO/outputs/hallu_v5_2_split.json}"
SIGNALS="${SIGNALS:-$REPO/outputs/latticespike_v5_2_signals_L4}"
LAYERS="${LAYERS:-4,8,12,16}"
MODEL="meta-llama/Llama-3.2-1B-Instruct"
REVISION="9213176726f574b556790deb65791e0c5aa438b6"
ROOT="$REPO/outputs/multilayer"
EPOCHS=10

mkdir -p "$ROOT"
echo "### MULTI-LAYER EXPERIMENT — layers ${LAYERS}, ${EPOCHS} epoch cap, 3 seeds"
echo "signals: $SIGNALS   (first run extracts; later seeds reuse)"
echo "started: $(date -Is)"; echo

# Objective identical to the current best, so the ONLY difference from the
# 5090 baseline is which hidden layers feed the box encoder.
BEST="--lam-contain 0 --lam-vol 0 --lam-recon 0 --lam-mrl 0 \
      --rank-objective ap --lam-rank 12 --ap-tau 0.25"

STATUS_FILE="$ROOT/status.tsv"; : > "$STATUS_FILE"

# Which seeds to run. Default is seed 42 alone: it answers the only question that
# matters first -- does multi-layer beat the single-layer baseline at all -- for
# roughly half the wall clock. Seeds 43/44 only convert a signal into a confirmed
# one, and if 42 is flat there is nothing to confirm.
#   SEEDS="42 43 44" scripts/run_multilayer_aws.sh   # full three-seed run
#
# Sequential on purpose. Concurrent processes would each try to extract the same
# 9,817 records at once on a 23 GB card.
for SEED in ${SEEDS:-42}; do
  name="ml_s${SEED}"
  out="$ROOT/$name"; log="$ROOT/$name.log"
  mkdir -p "$out"
  echo "=== [$name] start $(date -Is)"
  start=$(date +%s)
  # shellcheck disable=SC2086
  "$PY" -u scripts/train_latticespike_llm_eval.py \
      --input "$DATA" --split-summary "$SPLIT" --signals-dir "$SIGNALS" \
      --model "$MODEL" --model-revision "$REVISION" \
      --allow-unverified-replay --tie-box-token-emb \
      --box-hidden-layers "$LAYERS" \
      $BEST --epochs "$EPOCHS" --seed "$SEED" \
      --out-dir "$out" --summary "$out/summary.json" \
      > "$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_FILE"
  else
    echo "=== [$name] FAILED (rc=$rc) after ${elapsed}s"
    tail -n 20 "$log" | sed "s/^/    [$name] /"
    printf '%s\tFAILED(rc=%s)\t%s\n' "$name" "$rc" "$elapsed" >> "$STATUS_FILE"
  fi
  echo
done

echo "=== multi-layer finished $(date -Is) ==="; sort "$STATUS_FILE"; echo
"$PY" scripts/collect_variant_results.py "$ROOT" | tee "$ROOT/results_table.txt"
echo
echo "Single-layer baseline (5090, 3 seeds): VAL 0.2683 +/- 0.0088, TEST 0.2741 +/- 0.0238"
echo "Learned layer-mixture weights are stored in each checkpoint under"
echo "model.layer_mixer.weights -- read them to see which layers the model used."