#!/usr/bin/env bash
# =============================================================================
# run_len_sweep.sh -- locate the context length at which generative retrieval
# breaks down.
#
# WHY: at PROMPT_LEN=16360 every arm scores 0% even at depth 1.0, where the
# needle sits ~30 tokens before the question -- yet the same harness and format
# HIT at 2048. So retrieval degrades with TOTAL CONTEXT LENGTH, not with
# needle-to-question distance, and 16k alone cannot discriminate the arms.
#
# CRITICAL: T stays fixed at 16384 across every length. T sets the YaRN cos/sin
# cache (built for seq_len=T at scale 4), so varying it would change the model's
# positional encoding and confound length with RoPE. Fixing T and varying only
# PROMPT_LEN isolates context length as the sole variable.
#
# Cost is dominated by the longest point, which the main run already covers, so
# the added lengths are cheap (~5 min wall on 4 GPUs).
# =============================================================================
set -uo pipefail
cd "$(dirname "$0")"

PY=${PY:-/root/Adhi/BluTrain/envv/bin/python}
LENS=${LENS:-"1024 2048 4096 8192"}
TASKS=${TASKS:-"niah_single val_backward"}
CKPTS=${CKPTS:-"base arm_a arm_b arm_c"}
N=${N:-50}
MAX_NEW=${MAX_NEW:-24}
DEPTHS=${DEPTHS:-"0.0,0.25,0.5,0.75,1.0"}
GPUS=${GPUS:-0,1,2,3}

for L in $LENS; do
  echo "=============== PROMPT_LEN=$L (T fixed at 16384) ==============="
  PY="$PY" OUT="longeval_len_${L}" TASKS="$TASKS" CKPTS="$CKPTS" N="$N" \
  PROMPT_LEN="$L" MAX_NEW="$MAX_NEW" T=16384 DEPTHS="$DEPTHS" \
  GPUS="$GPUS" NP=1 BF16=1 GREEDY=1 \
    ./run_longeval.sh 2>&1 | grep -aE '^\[|accuracy:|depth|PRIMARY|^  [0-9]' | tail -25
done

echo
echo "=============== length sweep summary ==============="
$PY - "$LENS" "$CKPTS" <<'PY'
import csv, os, sys
lens = sys.argv[1].split(); models = sys.argv[2].split()
print("  depth-1.0 accuracy % (needle ~30 tok before the question = easiest case)")
print("  " + "len".ljust(8) + "".join(m.rjust(10) for m in models))
for L in lens:
    p = f"longeval_len_{L}/depth_summary.csv"
    row = {}
    if os.path.exists(p):
        for r in csv.DictReader(open(p)):
            if abs(float(r["depth"]) - 1.0) < 1e-9:
                row[r["model"]] = float(r["pct"])
    print("  " + L.ljust(8) + "".join(
        (f"{row[m]:9.1f}" if m in row else "        -").rjust(10) for m in models))
PY