#!/usr/bin/env bash
# Controlled A/B driver for the slow-mode readout (docs/SLOW_MODE_READOUT.md), derived
# from scripts/train_leakspec_arm.sh (same recipe, seed, stream skip, chunking).
# Extra per-chunk measurements: 2048-transition val (eval_windows.py) and the slow-mode
# probes (probe_leak_spectrum.py --slow-probes).  Artifacts: docs/artifacts/slow_mode_readout.
#
# One ARM = the exact fieldsat recipe (scripts/train_fieldsat_branch.sh: live recipe +
# --last-atom-dropout 0.5 + LR 1.5e-5 / SLR 3.6e-5, seed 20260913, stream-skip = START-496000)
# resumed from START_CKPT, in CHUNK-step chunks up to END_TARGET, with its own output dir and
# checkpoint suffix.  The ONLY difference between arms is LEAK_FLAGS:
#   ARM=A  LEAK_FLAGS="--field-leak 0.3"                                  (control)
#   ARM=B  LEAK_FLAGS="--field-leak-spectrum ... --slow-mode-readout --slow-memory-loss W ..."
# Resumable: restarts from the newest ckpt of this arm (<= END_TARGET), else START_CKPT.
# After each chunk: val normal / field-alone on the fixed 4876k ring, probe_leak_spectrum
# (cap fraction, a_hat linear next-byte CE, prompt-memory @16/32/64), coherence chat
# (3 prompts, standard params), seeded chat (seeds 1-3 x T 0.8/0.6/1.0) + prompt-dependence;
# output-dependence (JS divergence of next-byte distributions across prompts at the reply
# start); one row per chunk in $TSV (docs/artifacts/field_leak_spectrum/tsv_row.py).
#
# Env: ARM OUT SUF LEAK_FLAGS TSV [START_CKPT START_STEP END_TARGET CHUNK THREADS MEASURE_THREADS LR SLR LADROP]
#   THREADS=8          OMP/MKL threads for training
#   MEASURE_THREADS=1  threads for the (sequential, tiny-op) measurement tools; 8-thread OMP
#                      measurement next to an 8-thread training run thrashes (seen: chat_seeds
#                      35 CPU-min, training 0.3 steps/s), and thread count never touches training.
set -u
cd /workspace/exports/atom-ai
ARM=${ARM:?ARM}
OUT=${OUT:?OUT}
SUF=${SUF:?SUF}
LEAK_FLAGS=${LEAK_FLAGS:?LEAK_FLAGS}
TSV=${TSV:?TSV}
START_CKPT=${START_CKPT:-checkpoints/byte_tick_fieldsat/atom_native_step_4996000_d32_ms1M_lrd_b4828_fieldsat.pt}
START_STEP=${START_STEP:-4996000}
END_TARGET=${END_TARGET:-5004000}
CHUNK=${CHUNK:-8000}
THREADS=${THREADS:-8}
MEASURE_THREADS=${MEASURE_THREADS:-1}
LR=${LR:-1.5e-5}
SLR=${SLR:-3.6e-5}
LADROP=${LADROP:-0.5}
BASE=496000
A=docs/artifacts/slow_mode_readout
AB=$A/arm_$ARM
LOGD=logs/slowread/arm_$ARM
RING=checkpoints/byte_tick_fieldsat/val_ring_4876000.pt
PY=.venv/bin/python
export PYTHONPATH=. OMP_NUM_THREADS=$THREADS MKL_NUM_THREADS=$THREADS OPENBLAS_NUM_THREADS=$THREADS
mkdir -p "$OUT" "$AB" "$LOGD"
ts() { date -Is; }
exec 9>"$LOGD/lock"
if ! flock -n 9; then echo "$(ts) arm $ARM already running; exiting"; exit 1; fi
echo "=== $(ts) slowread arm=$ARM OUT=$OUT SUF=$SUF LEAK_FLAGS=[$LEAK_FLAGS] START=$START_STEP END=$END_TARGET THREADS=$THREADS ==="

latest() {
  ls "$OUT"/atom_native_step_*_${SUF}.pt 2>/dev/null | sed -n 's/.*step_\([0-9]*\)_.*/\1/p' | sort -n \
    | awk -v lo="$START_STEP" -v hi="$END_TARGET" '$1>lo && $1<=hi' | tail -1
}
has_row() { [ -f "$TSV" ] && awk -F'\t' -v s="$1" -v a="$ARM" 'NR>1 && $1==a && $2==s {f=1} END {exit !f}' "$TSV"; }

measure() {
  local END=$1 CK="$OUT/atom_native_step_$1_${SUF}.pt" ML="$LOGD/measure_$1.log"
  echo "$(ts) measure $ARM $END start (threads=$MEASURE_THREADS)"
  local OMP_NUM_THREADS=$MEASURE_THREADS MKL_NUM_THREADS=$MEASURE_THREADS OPENBLAS_NUM_THREADS=$MEASURE_THREADS
  export OMP_NUM_THREADS MKL_NUM_THREADS OPENBLAS_NUM_THREADS
  "$PY" -u tools/eval_last_atom_ablation.py --end-steps 4876000 --ring-cache "$RING" \
    --checkpoints "$CK" --out "$AB/val_ablation_$END.json" >> "$ML" 2>&1
  "$PY" -u docs/artifacts/field_leak_spectrum/eval_windows.py "$AB/eval_windows_$END.json" "$CK" >> "$ML" 2>&1
  "$PY" -u tools/probe_leak_spectrum.py --checkpoint "$CK" --variants ckpt --slow-probes \
    --out "$AB/probe_$END.json" >> "$ML" 2>&1
  local G="$AB/gen_coherent_$END.txt"; : > "$G"
  for p in "Bonjour, comment ça va ?" "Qui es-tu ?" "Il était une fois"; do
    RESP=$("$PY" tools/chat_atom_native.py --checkpoint "$CK" --prompt "$p" --no-role-prime \
      --max-packets 96 --max-length 100 --temperature 0.8 --top-k 12 2>/dev/null | sed -n 's/^Response: //p' | head -1)
    echo "PROMPT: $p" >> "$G"; echo "RESP: $RESP" >> "$G"
  done
  "$PY" tools/chat_seeds_probe.py --checkpoint "$CK" --seeds 1 2 3 --temps 0.8 0.6 1.0 \
    --out "$AB/chat_seeds_$END.json" > "$AB/chat_seeds_$END.txt" 2>&1
  "$PY" -u tools/probe_leak_spectrum.py --checkpoint "$CK" --variants ckpt --skip-episode --skip-memory \
    --output-dependence --out "$AB/outdep_$END.json" >> "$ML" 2>&1
  "$PY" "$A/tsv_row.py" "$ARM" "$END" "$CK" "$TSV" "$AB" "$OUT/run_metrics_${END}_${SUF}.json"
  echo "$(ts) measure $ARM $END done"
}

recover_metrics() {
  # ckpt produced by a run whose driver shell was stopped: recover its metrics files
  # (run_metrics.json in $OUT belongs to this ckpt iff its checkpoint path matches).
  local END=$1 NAME="atom_native_step_$1_${SUF}.pt"
  if [ ! -f "$AB/run_metrics_${END}.json" ] && grep -q "$NAME" "$OUT/run_metrics.json" 2>/dev/null; then
    cp "$OUT/run_metrics.json" "$OUT/run_metrics_${END}_${SUF}.json"
    cp "$OUT/run_metrics.json" "$AB/run_metrics_${END}.json"
    cp "$OUT/loss_trajectory.json" "$AB/loss_trajectory_${END}.json" 2>/dev/null || true
    echo "$(ts) recovered metrics for existing $NAME"
  fi
}

CUR=$(latest); [ -z "$CUR" ] && CUR=$START_STEP
while [ "$CUR" -lt "$END_TARGET" ]; do
  END=$((CUR + CHUNK)); [ "$END" -gt "$END_TARGET" ] && END=$END_TARGET
  STEPS=$((END - CUR)); SKIP=$((CUR - BASE))
  if [ "$CUR" -eq "$START_STEP" ]; then RESUME="$START_CKPT"; else RESUME="$OUT/atom_native_step_${CUR}_${SUF}.pt"; fi
  NAME="atom_native_step_${END}_${SUF}.pt"
  LOG="$LOGD/train_${CUR}_${END}.log"
  if [ ! -f "$OUT/$NAME" ]; then
    echo "$(ts) chunk $CUR->$END skip=$SKIP resume=$RESUME log=$LOG"
    T0=$(date +%s)
    # shellcheck disable=SC2086
    "$PY" -u tools/run_atom_native.py \
      --stream --data-glob 'data/corpus_fr_medium*.txt' --loop-shards --chunk-bytes 65536 \
      --stream-skip-packets "$SKIP" --resume "$RESUME" --output-dir "$OUT" \
      --checkpoint-name "$NAME" --steps "$STEPS" \
      --d-model 32 --n-modes 32 --n-atoms-max 64 --max-span-bytes 1 \
      --field-obligatory-hard --no-enable-merge --no-payload-copy --last-atom-readout \
      --efference-every 10 --learning-rate "$LR" --surface-learning-rate "$SLR" --seed 20260913 \
      --atom-flush-every 64 --episode-length 8000 --episode-reset --slow-every 1 \
      --field-loss-weight 0.05 --field-contrast-weight 0 --allow-parallel-train \
      --last-atom-dropout "$LADROP" $LEAK_FLAGS \
      > "$LOG" 2>&1
    RC=$?
    if [ "$RC" -ne 0 ] || [ ! -f "$OUT/$NAME" ]; then echo "$(ts) chunk FAILED rc=$RC (see $LOG)"; exit 1; fi
    cp "$OUT/run_metrics.json" "$OUT/run_metrics_${END}_${SUF}.json"
    cp "$OUT/run_metrics.json" "$AB/run_metrics_${END}.json"
    cp "$OUT/loss_trajectory.json" "$AB/loss_trajectory_${END}.json" 2>/dev/null || true
    echo "$(ts) chunk $CUR->$END saved $NAME wall=$(( $(date +%s) - T0 ))s"
  else
    recover_metrics "$END"
  fi
  has_row "$END" || measure "$END"
  CUR=$END
done
# measure any checkpoint of this arm in (START_STEP, END_TARGET] that has no TSV row yet
for E in $(ls "$OUT"/atom_native_step_*_${SUF}.pt 2>/dev/null | sed -n 's/.*step_\([0-9]*\)_.*/\1/p' | sort -n); do
  if [ "$E" -gt "$START_STEP" ] && [ "$E" -le "$END_TARGET" ] && ! has_row "$E"; then
    recover_metrics "$E"; measure "$E"
  fi
done
echo "$(ts) DONE arm $ARM at $CUR"