#!/usr/bin/env bash
# =============================================================================
# needle_sweep.sh -- run the retrieval probe across an extension-factor row set
# for ONE model, then score it. Forward-only, cache-only (no fine-tuning).
#
# Reuses the existing CP_EVAL_PPL + CP_EVAL_TOKENS_BIN path unchanged: the .bin
# holds synthetic trials, and per-position mean_nll in the CSV gives retrieval
# NLL at the answer span. See needle_gen.py for the fixed-slot construction.
#
# GROUP-SLOT FAN-OUT: slots = NGPU / CP_SIZE, so EVERY row parallelises --
#   T<=128k -> CP=1 -> 8 concurrent   |  256k -> CP=2 -> 4   |  512k -> CP=4 -> 2
# Each group is pinned to a contiguous device subset so cudaSetDevice(rank) binds
# rank->GPU correctly. Concurrency changes only WHEN work happens, never what is
# computed (independent processes, separate CSVs), so no determinism argument is
# needed -- unlike the search's EvaluatorPool.
#
# RESUMABLE: one marker per (T, cache, variant, block) in $STATE_DIR.
#
# Launch (114M 4k base, NTK, all five factors):
#   MODEL_KEY=114m_run1 CKPT_RUN=1 \
#   MODEL_ARCH="CP_N_EMBD=768 CP_N_LAYER=12 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=2048 CP_WEIGHT_TYING=1" \
#   FACTORS="8 16 32" CACHES="ntk" \
#     nohup bash Tests/bluscriptcp/needle_sweep.sh > needle_114m.log 2>&1 &
#
# The single searched cell (48M@16x) is a separate invocation:
#   MODEL_KEY=48m_run200 CKPT_RUN=200 MODEL_ARCH="..." FACTORS="16" \
#   CACHES="searched" SEARCHED_LAMBDA=longrope_best_48m_s16_repr.txt bash ...
# =============================================================================
set -euo pipefail
cd "$(dirname "$0")/../.."   # -> repo root

MODEL_KEY=${MODEL_KEY:?set MODEL_KEY (label for the results rows, e.g. 114m_run1)}
# INIT_FROM=<path/to/base.ckpt> loads WEIGHTS ONLY instead of resuming a run.
# Needed whenever the checkpoint's optimizer state cannot be read back in the
# EVAL configuration: a ZeRO/Muon-trained base evaluated forward-only builds a
# plain nn::AdamW, which would try to parse ShardedMuon's state and fail. Also
# the only way to evaluate a base trained at a different world size, since ZeRO
# state is sharded to the size that wrote it. A forward-only probe needs no
# optimizer state at all, so this is the correct loader for it.
INIT_FROM=${INIT_FROM:-}
if [ -n "$INIT_FROM" ]; then
  [ -f "$INIT_FROM" ] || { echo "INIT_FROM: no such file: $INIT_FROM" >&2; exit 1; }
  CKPT_RUN=${CKPT_RUN:-0}
  CKPT_ENV="CP_INIT_FROM=$INIT_FROM"
else
  CKPT_RUN=${CKPT_RUN:?set CKPT_RUN (or INIT_FROM=<path>)}
  CKPT_ENV="CP_CKPT_RESUME=$CKPT_RUN"
fi
MODEL_ARCH=${MODEL_ARCH:?set MODEL_ARCH (must match how CKPT_RUN was trained)}

FACTORS=${FACTORS:-"8 16 32 64 128"}     # extension factors s; T = PHYS_T * s
CACHES=${CACHES:-"ntk"}                  # ntk | yarn | pi | searched
# YARN_SCALE: evaluate under the trainer's YaRN cache instead of the LongRoPE lambda
# ramp. These are DIFFERENT ways to reach the same T and a checkpoint must be scored
# under the one it was TRAINED with, or the numbers are meaningless -- measured:
# run7 (extended at YARN_SCALE=4, m=1.1386) scored under the lambda ramp (scale=1.0)
# had its retrieval FLOOR collapse from 4.17 to 11.34 nats, i.e. it failed even with
# the needle 12 tokens from the probe. bluscriptCP.cpp:914 states the intended matrix:
#   base ckpt     + YARN_SCALE=1   (naive extrapolation @ T)
#   extended ckpt + YARN_SCALE=4   (YaRN @ T)
# Set YARN_SCALE to take that route; leave it empty for the LongRoPE route.
YARN_SCALE=${YARN_SCALE:-}
SEARCHED_LAMBDA=${SEARCHED_LAMBDA:-}     # required when CACHES contains "searched"
PHYS_T=${PHYS_T:-4096}
HEAD_DIM=${HEAD_DIM:-64}
ROPE_THETA=${ROPE_THETA:-500000}
BLOCKS=${BLOCKS:-5}
NEEDLE_LEN=${NEEDLE_LEN:-8}
SEED=${SEED:-0}
GPUS=${GPUS:-0,1,2,3,4,5,6,7}
FILLER=${FILLER:-Data_Loader/Data/edufineweb_val_00000.bin}
DATA_ROOT=${DATA_ROOT:-$PWD/Data_Loader/Data}
NEEDLE_DIR=${NEEDLE_DIR:-needles}
# Absolute offset for the probe+answer span (needle_gen --probe-pos). Unset = the
# end of the trial, the historical layout. Set it to decouple needle-to-probe
# DISTANCE from the needle's ABSOLUTE POSITION, which the default layout welds
# into a single variable. Filenames do NOT encode it, so give each probe position
# its own NEEDLE_DIR/CSV_DIR/STATE_DIR.
PROBE_POS=${PROBE_POS:-}
CSV_DIR=${CSV_DIR:-needle_csv}
STATE_DIR=${STATE_DIR:-needle_state}
RESULTS=${RESULTS:-needle_results.csv}
# Per-rank forward-only logits are (T/CP)*vocab*4 bytes. Keep <= this so a 49 GB
# card holds logits + activations. Pre-flight 2b must confirm it (and that logits
# are fp32 -- a bf16 model triggers an as_type copy that doubles the peak).
MAX_TOK_PER_RANK=${MAX_TOK_PER_RANK:-131072}

EXE=./build/bluscriptCP_exec
export LD_LIBRARY_PATH=BluTrain/Tensor-Implementations/lib:BluTrain/Profiler/lib:${LD_LIBRARY_PATH:-}
IFS=',' read -r -a GPUARR <<< "$GPUS"
NGPU=${#GPUARR[@]}
KV=$(sed -E 's/.*CP_N_KVHEAD=([0-9]+).*/\1/' <<< "$MODEL_ARCH")
mkdir -p "$NEEDLE_DIR" "$CSV_DIR" "$STATE_DIR"
log(){ echo "[$(date '+%F %H:%M:%S')] $*"; }

# ---- pre-flight -------------------------------------------------------------
[ -x "$EXE" ] || { echo "ERROR: $EXE missing -- build first"; exit 1; }
[ -f "$FILLER" ] || { echo "ERROR: filler stream $FILLER not found"; exit 1; }
# Detect a live evaluator by RESOLVED EXECUTABLE, not by command-line text. `pgrep -f
# bluscriptCP_exec` also matches any shell whose argv merely CONTAINS that string --
# including the wrapper that launches this script -- which produced a false "already
# running" abort. `pgrep -x` is no good either: comm truncates to 15 chars and
# "bluscriptCP_exec" is 16. Comparing /proc/PID/exe is exact.
exec_running(){
  local want; want=$(readlink -f "$EXE") || return 1
  local p
  for p in /proc/[0-9]*; do
    [ "$(readlink -f "$p/exe" 2>/dev/null)" = "$want" ] && { echo "${p#/proc/}"; return 0; }
  done
  return 1
}
if pids=$(exec_running); then
  # ALLOW_CONCURRENT_EXE=1 permits a DELIBERATE overlap -- e.g. probing a mid-training
  # checkpoint while the pretrain keeps running. Only use it when the memory headroom is
  # known: a forward-only probe is ~10-15 GB against a training job's ~47 GB on an 80 GB
  # card. The context_rot_2026-08-19 copies of this script already carried this flag; the
  # tracked one never got it.
  if [ "${ALLOW_CONCURRENT_EXE:-0}" = "1" ]; then
    log "WARNING: $EXE already running (pid $pids); ALLOW_CONCURRENT_EXE=1 -- proceeding"
  else
    echo "ERROR: $EXE already running (pid $pids) -- this sweep would oversubscribe GPUs $GPUS"
    echo "  set ALLOW_CONCURRENT_EXE=1 to override when the overlap is intended"
    exit 1
  fi
fi
case " $CACHES " in *" searched "*)
  [ -n "$SEARCHED_LAMBDA" ] || { echo "ERROR: CACHES has 'searched' but SEARCHED_LAMBDA is unset"; exit 1; }
  [ -f "$SEARCHED_LAMBDA" ] || { echo "ERROR: SEARCHED_LAMBDA '$SEARCHED_LAMBDA' not found"; exit 1; };;
esac

# CP degree per row, from the plan's CP table (memory-driven, not a free choice).
cp_for_T(){
  local T=$1 cp=1
  while [ $(( T / cp )) -gt "$MAX_TOK_PER_RANK" ]; do cp=$(( cp * 2 )); done
  echo "$cp"
}
# Verified topology only: pure ring>=4 is known-broken in this repo.
topo_for(){
  local cp=$1
  if   [ "$cp" = "1" ]; then echo ""
  elif [ $(( KV % cp )) -eq 0 ]; then echo "CP_ATTN_MODE=ulysses"
  elif [ $(( cp % KV )) -eq 0 ]; then echo "CP_ATTN_MODE=hybrid CP_ULYSSES_SIZE=$KV"
  else echo "UNSUPPORTED"; fi
}

# lambda file per (cache, head_dim, theta, s). NTK/PI are theta-independent closed
# forms; YaRN depends on base via find_correction_range, so it MUST be regenerated
# per theta. All are head_dim-dependent (vector length = head_dim/2).
lambda_file(){
  local cache=$1 s=$2
  if [ "$cache" = "searched" ]; then echo "$SEARCHED_LAMBDA"; return; fi
  local out="$NEEDLE_DIR/lam_${cache}_hd${HEAD_DIM}_th${ROPE_THETA}_s${s}.txt"
  if [ ! -f "$out" ]; then
    python3 - "$cache" "$s" "$out" "$HEAD_DIM" "$ROPE_THETA" "$PHYS_T" <<'PY'
import sys, importlib.util
cache, s, out, hd, theta, orig = (sys.argv[1], int(sys.argv[2]), sys.argv[3],
                                  int(sys.argv[4]), float(sys.argv[5]), float(sys.argv[6]))
spec = importlib.util.spec_from_file_location('lrs', 'Tests/bluscriptcp/longrope_search.py')
lrs = importlib.util.module_from_spec(spec); spec.loader.exec_module(lrs)
half, lam_max = hd // 2, s * 1.25
if   cache == 'ntk':  lam = lrs.seed_NTK(hd, s, lam_max)
elif cache == 'pi':   lam = lrs.seed_PI(half, s, lam_max)
elif cache == 'yarn': lam = lrs.seed_YaRN(hd, s, orig, 32.0, 1.0, theta, lam_max)
else: sys.exit(f"unknown cache '{cache}'")
lrs.write_factors(out, lam, 0, s, int(orig * s))
PY
  fi
  echo "$out"
}

# Overridable so a REFINED distance grid can be added around a crossing. Set
# VARIANTS to the UNION of old and new points: already-completed ones are skipped
# by their state markers, so only the new points actually run, and the scorer then
# reports the whole grid in one table.
VARIANTS=${VARIANTS:-"p0 p0.25 p0.5 p0.75 p1 absent null"}
log "needle sweep: model=$MODEL_KEY run=$CKPT_RUN factors=[$FACTORS] caches=[$CACHES] blocks=$BLOCKS"
log "  GPUS=$GPUS (n=$NGPU) kv=$KV head_dim=$HEAD_DIM theta=$ROPE_THETA"

for S in $FACTORS; do
  T=$(( PHYS_T * S ))
  CP=$(cp_for_T "$T")
  TOPO=$(topo_for "$CP")
  [ "$TOPO" = "UNSUPPORTED" ] && { echo "ERROR: no verified topology for CP=$CP with kv=$KV at T=$T"; exit 1; }
  [ "$CP" -le "$NGPU" ] || { echo "ERROR: T=$T needs CP=$CP but only $NGPU GPU(s) given"; exit 1; }
  N=$([ "$T" -le 131072 ] && echo 8 || echo 4)      # pinned trials-per-block (plan A1)
  SLOTS=$(( NGPU / CP ))
  log "== s=$S T=$T CP=$CP ${TOPO:-(single-GPU)} N=$N slots=$SLOTS =="

  # .bin files depend only on (T, variant, block) -- shared across models and caches.
  for V in $VARIANTS; do for B in $(seq 0 $((BLOCKS-1))); do
    [ -f "$NEEDLE_DIR/needle_T${T}_${V}_b${B}.bin" ] && continue
    python3 Tests/bluscriptcp/needle_gen.py --T "$T" --variant "$V" --block "$B" \
        --trials "$N" --needle-len "$NEEDLE_LEN" --filler "$FILLER" \
        --seed "$SEED" --out-dir "$NEEDLE_DIR" \
        ${PROBE_POS:+--probe-pos "$PROBE_POS"} >/dev/null
  done; done

  for CACHE in $CACHES; do
    LAM=$(lambda_file "$CACHE" "$S")
    # The lambda vector is length-specific: reject a mismatch loudly rather than
    # scoring a cache tuned for another length (the T == S_search invariant).
    SS=$(awk '/^S_search/{print $2}' "$LAM")
    [ "$SS" = "$T" ] || { echo "ERROR: $LAM has S_search=$SS but this row is T=$T"; exit 1; }

    # ---- job list, then a group-slot loop ----
    jobs=()
    for V in $VARIANTS; do for B in $(seq 0 $((BLOCKS-1))); do jobs+=("$V:$B"); done; done
    running=0; g=0
    for job in "${jobs[@]}"; do
      V=${job%%:*}; B=${job##*:}
      MK="$STATE_DIR/${MODEL_KEY}.T${T}.${CACHE}.${V}.b${B}.done"
      OUTCSV="$CSV_DIR/ppl_needle_T${T}_${V}_b${B}.csv"
      [ -f "$MK" ] && continue
      # pin this group to its own contiguous device subset
      devs=$(IFS=,; echo "${GPUARR[*]:$((g*CP)):$CP}")
      if [ -n "$YARN_SCALE" ]; then
        ROPE_ENV="YARN_SCALE=$YARN_SCALE"          # trainer's YaRN cache
      else
        ROPE_ENV="CP_LONGROPE_FACTORS=$LAM"        # LongRoPE lambda ramp
      fi
      (
        env $MODEL_ARCH CP_SIZE=$CP $TOPO \
            CP_EVAL_PPL=1 CP_T=$T CP_B=1 \
            CP_EVAL_WINDOWS=$N CP_EVAL_SKIP_WINDOWS=0 \
            CP_EVAL_TOKENS_BIN="$NEEDLE_DIR/needle_T${T}_${V}_b${B}.bin" \
            $ROPE_ENV $CKPT_ENV \
            CP_EVAL_OUT="$OUTCSV" CP_DATA_ROOT="$DATA_ROOT" \
            CUDA_VISIBLE_DEVICES=$devs mpirun -np $CP "$EXE" \
          > "$CSV_DIR/${MODEL_KEY}_T${T}_${CACHE}_${V}_b${B}.log" 2>&1 \
          && touch "$MK"
      ) &
      g=$(( (g + 1) % SLOTS )); running=$(( running + 1 ))
      if [ "$running" -ge "$SLOTS" ]; then wait -n; running=$(( running - 1 )); fi
    done
    wait

    # every job must have produced its marker, or scoring would read a stale CSV
    miss=0
    for V in $VARIANTS; do for B in $(seq 0 $((BLOCKS-1))); do
      [ -f "$STATE_DIR/${MODEL_KEY}.T${T}.${CACHE}.${V}.b${B}.done" ] || {
        echo "  MISSING: $V b$B -- see $CSV_DIR/${MODEL_KEY}_T${T}_${CACHE}_${V}_b${B}.log"; miss=1; }
    done; done
    [ "$miss" = "0" ] || { echo "ERROR: s=$S cache=$CACHE incomplete; not scoring"; exit 1; }

    log "  scoring s=$S cache=$CACHE -> $RESULTS"
    python3 Tests/bluscriptcp/needle_score.py \
        --needle-dir "$NEEDLE_DIR" --csv-dir "$CSV_DIR" --T "$T" --blocks "$BLOCKS" \
        --model "$MODEL_KEY" --cache "$CACHE" --factor "$S" \
        --head-dim "$HEAD_DIM" --rope-theta "$ROPE_THETA" --out "$RESULTS" \
        --p-fracs "$(tr ' ' '\n' <<< "$VARIANTS" | grep '^p' | paste -sd,)"
  done
done

log "SWEEP DONE -> $RESULTS  (CSVs in $CSV_DIR, markers in $STATE_DIR)"