#!/usr/bin/env bash
# 1. x64 sweep -- does the scale curve keep rising at 64k?
# 2. val_backward -- the DISCRIMINATING task. 2026-09-09 found niah_single is
#    non-discriminating (Qwen3 0.6B-8B all 100%, zero variation over 13x params)
#    while val_backward separates them and ours were at chance. If today's NIAH
#    gains are real capability, val_backward should move too. If they are a
#    compression artefact, it will not.
set -uo pipefail
cd "$(dirname "$0")"
export PATH=/root/Adhi/BluTrain/envv/bin:$PATH
source /root/Adhi/BluTrain/envv/bin/activate 2>/dev/null
L=final_logs; mkdir -p "$L"; M="$L/master.log"
say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$M" >&2; }
HAY=Data_Loader/Data/b2_8192_16384/val_shard_00000.bin
[ "$(stat -c%s "$HAY")" = 199754694 ] || { say "FATAL haystack differs"; exit 1; }
A680="CP_N_EMBD=1536 CP_N_LAYER=24 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=1"

go(){ # T SCALE PL TASKS CKPTS EXTRA ARCH tag
  local T=$1 S=$2 PL=$3 TK=$4 CKS=$5 EX=$6 AR=$7 TAG=$8
  local O="fin_${TAG}_${PL}"
  [ $((S*4096)) -eq $T ] || { say "YaRN invariant $S*4096 != $T"; return 1; }
  [ -s "$O/summary.csv" ] && [ "$(wc -l < "$O/summary.csv")" -gt 1 ] && {
    say "  $TAG $PL cached: $(tail -n +2 $O/summary.csv|tr '\n' ' ')"; return 0; }
  say "  $TAG len $PL (T=$T x$S -> virtual $((PL/S))) tasks=$TK"
  A1B="$AR" HAY="$HAY" GEN_DATA_ROOT="Data_Loader/Data/evalsrc" \
  CKPT_EXTRA="$EX" CKPTS="$CKS" TASKS="$TK" N=100 \
  T=$T YARN=$S YARN_ORIG=4096 PROMPT_LEN=$PL GPUS=0,1 NP=1 OUT="$O" \
    ./run_longeval.sh >> "$L/${TAG}_${PL}.log" 2>&1
  if [ -s "$O/summary.csv" ] && [ "$(wc -l < "$O/summary.csv")" -gt 1 ]; then
    say "    $(head -1 $O/summary.csv)"; tail -n +2 "$O/summary.csv" | sed 's/^/      /' | tee -a "$M" >&2
  else say "    FAILED"; grep -ahA2 'what()' "$O"/gen/*.log 2>/dev/null|head -3|tee -a "$M" >&2; fi
}
E680="p1final=blumodelcp_run12_step_1200.ckpt base680=varybatch.ckpt"

say "=== PART 1: x64 sweep (T=262144) -- does scale still help at 64k? ==="
for PL in 4096 32744 65512; do go 262144 64 $PL niah_single "p1final base680" "$E680" "$A680" x64; done

say "=== PART 2: val_backward at the best configs -- the discriminating task ==="
go  65536 16 16360 "val_backward niah_multikey" "p1final base680" "$E680" "$A680" vb16
go 131072 32 32744 "val_backward niah_multikey" "p1final base680" "$E680" "$A680" vb32
say "  and the 1.06B (trained to 64k) on the same task -- default 36-layer arch"
E106="mix1200=mix64k_step_1200.ckpt"
A106="CP_N_EMBD=1536 CP_N_LAYER=36 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=0"
go 131072 32 32744 "val_backward niah_multikey" "mix1200" "$E106" "$A106" vb106

say "=== ALL RESULTS ==="
for d in fin_*; do
  [ -f "$d/summary.csv" ] || continue
  say "  $d: $(head -1 $d/summary.csv) | $(tail -n +2 $d/summary.csv | tr '\n' ' ')"
done
say "DONE"