#!/usr/bin/env bash
# Probe-position sweep on ONE GPU. Two questions, one job set:
#   (1) ADJACENCY vs END-OF-DOCUMENT -- variant p1 (needle adjacent) is run at each
#       PROBE_POS. Arms B and C are ~4 nats worse than base at p1 when the probe sits
#       at the end (16364). If they recover with the probe mid-document, the defect is
#       end-of-document, not adjacency.
#   (2) U-SHAPE -- variant p<D> holds needle-probe DISTANCE fixed at 4087 tokens while
#       the probe's ABSOLUTE position moves. needle_gen's default (probe pinned at the
#       end) makes distance and position perfectly anti-correlated, so the main sweep
#       could not tell "far from the question" from "early in the document". This can.
set -uo pipefail
cd "$(dirname "$0")"
LOG=probepos_logs; mkdir -p $LOG
M=$LOG/master.log
say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a $M; }

CK=checkpoints_bluscriptcp/rank_0
declare -A MODELS=( [base]=blumodelcp_run7_step_750.ckpt \
                    [armB]=blumodelcp_run15_step_400.ckpt \
                    [armC]=blumodelcp_run16_step_400.ckpt )
# P=8192 first: the decisive middle-vs-end contrast. P=12288 adds the third curve point.
POSITIONS="8192 12288"
declare -A PFRAC=( [6144]=0.3331 [8192]=0.5001 [12288]=0.6670 )

say "probe-pos sweep starting on $(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)"
for P in $POSITIONS; do
  for MK in base armB armC; do
    OUT=probe_pos${P}_${MK}
    if [ -f "$OUT/reach.csv" ]; then say "  $OUT already done"; continue; fi
    say "  PROBE_POS=$P model=$MK  (variants: p${PFRAC[$P]} p1 null absent)"
    CKPT="$PWD/$CK/${MODELS[$MK]}" OUT=$OUT \
    PROBE_POS=$P VARIANTS="p${PFRAC[$P]} p1 absent null" \
    FACTORS=4 CACHES=yarn YARN_SCALE=4 GPUS=0 BLOCKS=3 \
      ./probe_1b.sh >>"$M" 2>&1 \
      && say "    -> $OUT/reach.csv" \
      || say "    WARN: $MK at P=$P returned nonzero (CSVs may still be usable)"
  done
done
say "ALL DONE"