#!/usr/bin/env bash
# Wait for build + download -> rebuild ablmix16 -> run the AdamW control arm.
set -uo pipefail; cd "$(dirname "$0")"
export PATH=/root/Adhi/BluTrain/envv/bin:$PATH
say(){ echo "[$(date -u +%H:%M:%S)] $*"; }
say "waiting for build..."
for i in $(seq 1 240); do
  [ -x build/bluscriptCP_exec ] && grep -q 'CP_FUSED_ROPE=1' build/objects/.cp_rope_flags 2>/dev/null && break
  grep -q 'BUILD FAILED\|NO libtensor' build.log 2>/dev/null && { say "BUILD FAILED"; tail -15 build.log; exit 1; }
  sleep 20
done
[ -x build/bluscriptCP_exec ] || { say "no binary"; tail -15 build.log; exit 1; }
say "binary $(stat -c%s build/bluscriptCP_exec) flags=$(cat build/objects/.cp_rope_flags)"
say "waiting for downloads..."
for i in $(seq 1 240); do
  grep -q 'DL_DONE' dl_abl.log 2>/dev/null && break; sleep 20; done
grep -q 'DL_DONE' dl_abl.log 2>/dev/null || { say "download incomplete"; tail -6 dl_abl.log; exit 1; }
say "downloads done"
if [ ! -d Data_Loader/Data/ablmix16 ]; then
  say "building ablmix16 (must match p1mix token-for-token)"
  python mix_buckets.py --out Data_Loader/Data/ablmix16 --data-root srcabl \
    --src b3_16384_32768=0.5 --src b2_8192_16384=0.3 --src b1_4096_8192=0.2 \
    --total-tokens 640e6 --shard-tokens 96e6 --window 16384 --grad-accum 32 2>&1 | tail -5
  cp -f Data_Loader/Data/b2_8192_16384/val_shard_00000.bin Data_Loader/Data/ablmix16/
fi
B3=$(grep -oE 'b3_16384_32768 target [0-9.]+ actual [0-9.]+ tokens [0-9]+' Data_Loader/Data/ablmix16/MANIFEST.txt | grep -oE '[0-9]+$')
[ "$B3" = "320011377" ] && say "mix VERIFIED token-identical to p1mix (b3=$B3)" || say "WARNING mix differs: b3=$B3 expected 320011377"
say "=== LAUNCHING: 8e-5 constant, AdamW only ==="
ARM=const GPU=0 OPT=adamw exec ./abl_phase1.sh