#!/usr/bin/env bash
# IN2 experiment driver: generate arm C, gate all arms, train A/B/C, probe each,
# read out, upload. Unattended, with hard gates -- every one of them exists because
# the failure it catches is SILENT and would otherwise produce a plausible-looking
# number for the wrong thing.
set -uo pipefail
cd "$(dirname "$0")"
ROOT="$(cd ../.. && pwd)"
export LD_LIBRARY_PATH="$ROOT/Tensor-Implementations/lib:$ROOT/Profiler/lib:${LD_LIBRARY_PATH:-}"
export OMPI_ALLOW_RUN_AS_ROOT=1 OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1
PY="$ROOT/envv/bin/python"; HF="$ROOT/envv/bin/hf"
LOGDIR=in2_logs; mkdir -p $LOGDIR
MASTER=$LOGDIR/driver_$(date +%Y%m%d_%H%M).log
say(){ echo "[$(date +%H:%M:%S)] $*" | tee -a "$MASTER"; }
die(){ say "FATAL: $*"; say "aborting -- see $MASTER"; exit 1; }

# ---- configuration -------------------------------------------------------
ARCH="CP_N_EMBD=1536 CP_N_LAYER=36 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=0"
BASE_CKPT=${BASE_CKPT:-$PWD/checkpoints_bluscriptcp/rank_0/blumodelcp_run7_step_750.ckpt}
DATA=${DATA:-/root/.cache/huggingface/hub/datasets--navingv--flux/snapshots/5232b3d133d2eb71905725f288ff6b56dc0df773}
T=16384; K=64; EXAMPLES=12816; SHARDS=3
# D_MIN is the generator's DEFAULT (64) and must stay there. It is tempting to set
# it to R so every arm-C needle lands in the failing region -- that is wrong, and the
# gate catches it: d_min feeds draw_sites(lo_margin=span+d_min), so raising it moves
# arm C's PROBES, not just its needles. At d_min=9835 arm C's answers were confined to
# the last ~6.5k of the document (6522 distinct offsets vs arm B's 16293) and cross-arm
# assertion 8 failed. Arm B can only absorb position effects while B and C share
# positions, which is exactly what raising d_min destroys -- and it re-creates the
# position-specific-circuit confound the interleaving exists to prevent.
# The consequence is accepted, not fixed: only ~8.2% of arm C's 820,224 triples land
# beyond R=9835 (67,653 of them). That fraction is the price of positional matching.
D_MIN=${D_MIN:-64}
R_MEASURED=${R_MEASURED:-9835}
FT_STEPS=${FT_STEPS:-400}; GBATCH=${GBATCH:-524288}; B=${B:-4}
YARN=4; GPUS=0,1,2,3; NP=4
HF_REPO=${HF_REPO:-unparallelled/BluScriptCP}
STAGE=/root/in2_upload; rm -rf $STAGE; mkdir -p $STAGE

# Single-instance lock. Three copies of this script ran concurrently once (one from
# each terminal), sharing in2_logs/, checkpoints_bluscriptcp/ and all four GPUs: 10
# exec processes contending, and they would have raced for run numbers and written
# each other's arm markers. A second launch must refuse, not collide.
LOCK=.in2_driver.lock
if [ -e "$LOCK" ] && kill -0 "$(cat "$LOCK" 2>/dev/null)" 2>/dev/null; then
  echo "REFUSING: another driver is already running (pid $(cat "$LOCK")). "
  echo "  If it is dead: rm $LOCK"
  exit 1
fi
echo $$ > "$LOCK"
trap 'rm -f "$LOCK"' EXIT INT TERM

say "IN2 driver starting.  base=$BASE_CKPT  d_min=$D_MIN  R=$R_MEASURED"
[ -x ./build/bluscriptCP_exec ] || die "no binary"
[ -f "$BASE_CKPT" ] || die "base checkpoint missing: $BASE_CKPT"
[ -x "$PY" ] || die "venv python missing"

# ---- 1. arm C data (A and B already generated) ---------------------------
if [ ! -f in2_data/arm_c/.done ]; then
  say "generating arm C (train) at --d-min $D_MIN (generator default; see note above)"
  FILL=$(ls Data_Loader/Data/flux_train_0000{1,2,3}*.bin 2>/dev/null | head -12 | tr '\n' ' ')
  [ -z "$FILL" ] && FILL=$(ls Data_Loader/Data/flux_train*.bin | head -12 | tr '\n' ' ')
  $PY Tests/bluscriptcp/in2_gen.py --arm c --T $T --examples $EXAMPLES --K $K --shards $SHARDS \
      --seed 0 --d-min $D_MIN --filler-shards $FILL --out-dir in2_data/arm_c >>"$MASTER" 2>&1 \
      || die "arm C train generation failed"
  say "generating arm C (val)"
  $PY Tests/bluscriptcp/in2_gen.py --arm c --T $T --examples 64 --split val --K $K --shards 1 \
      --seed 0 --d-min $D_MIN --filler-shards $(ls Data_Loader/Data/flux_val*.bin | head -1) \
      --out-dir in2_data/arm_c >>"$MASTER" 2>&1 || die "arm C val generation failed"
  touch in2_data/arm_c/.done
fi
grep -E 'distance:' "$MASTER" | tail -2 | while read -r l; do say "  arm C $l"; done

# ---- 2. GATE: 23 layout assertions, all arms, both splits ---------------
# This is what caught arm B silently dropping 213 of 1024 needles last time -- a
# desynchronised answer_start histogram destroys the B-vs-C matched-position
# discipline the whole primary claim rests on.
for SPLIT in train val; do
  say "gate ($SPLIT) on all three arms"
  OUT=$($PY Tests/bluscriptcp/in2_gate.py --dirs in2_data/arm_a in2_data/arm_b in2_data/arm_c \
        --split $SPLIT --B $B --dp 1 2>&1); echo "$OUT" >>"$MASTER"
  echo "$OUT" | grep -q "ALL PASS" || { echo "$OUT" | tail -20; die "GATE FAILED ($SPLIT)"; }
  say "  gate $SPLIT: ALL PASS"
done

# ---- 3. train the three arms --------------------------------------------
declare -A RUNNO
for ARM in a b c; do
  MK=$LOGDIR/arm_${ARM}.trained
  if [ -f "$MK" ]; then RUNNO[$ARM]=$(cat "$MK"); say "arm $ARM already trained (run ${RUNNO[$ARM]})"; continue; fi
  # One checkpoint step costs 17 GB across 4 ranks (3.9 GB model + 2 sidecars, per
  # rank). The first attempt filled a 193 GB disk: FREQ=100 x KEEP=2 x 3 arms is up
  # to 102 GB of arm checkpoints on top of the base. Saves then failed with a mere
  # [WARN] while training carried on, so arm C ran all 400 steps and reported "done"
  # with ZERO checkpoints on disk -- 45 minutes of GPU time for nothing. Check first.
  FREE_G=$(df -BG --output=avail / | tail -1 | tr -dc '0-9')
  say "  free disk: ${FREE_G}G"
  [ "${FREE_G:-0}" -ge 40 ] || die "only ${FREE_G}G free; one arm checkpoint needs ~17G plus a same-size .tmp during the write. Free space before continuing."

  # Prune sidecars whose .ckpt has already been rotated away. CP_CKPT_KEEP prunes
  # .ckpt files but NOT the zero_g* sidecars bluscriptCP writes alongside them, so
  # they accumulate silently -- 50 GB of them built up on the 600M run.
  for d in checkpoints_bluscriptcp/rank_*; do
    [ -d "$d" ] || continue
    K=$(ls $d/*.ckpt 2>/dev/null | grep -oE '_run[0-9]+_step_[0-9]+' | sort -u)
    for f in $d/*.zero_g*; do
      [ -e "$f" ] || continue
      S=$(echo "$f" | grep -oE '_run[0-9]+_step_[0-9]+')
      echo "$K" | grep -qx "$S" || rm -f "$f"
    done
    rm -f $d/*.tmp
  done

  L=$LOGDIR/arm_${ARM}_train.log
  say "training arm $ARM  ($FT_STEPS steps x $GBATCH tokens, T=$T, cp=$NP, YaRN=$YARN)"
  env CUDA_VISIBLE_DEVICES=$GPUS $ARCH \
    CP_INIT_FROM="$BASE_CKPT" \
    CP_SIZE=$NP CP_ATTN_MODE=ring CP_ATTN_FUSION=split \
    CP_T=$T YARN_SCALE=$YARN CP_B=$B CP_GLOBAL_BATCH=$GBATCH \
    CP_ZERO=1 CP_ZERO_V2=1 CP_ZERO_SYMMETRIC=1 CP_BF16=1 CP_OPTIMIZER=muon \
    CP_MAX_LR=1e-3 CP_MIN_LR=1e-4 CP_REWARMUP=100 CP_REWARMUP_PEAK=0.3 \
    CP_MAX_STEPS=$FT_STEPS CP_WARMUP=1 \
    CP_CKPT=1 CP_CKPT_FREQ=200 CP_CKPT_KEEP=1 CP_CKPT_NEW_RUN=1 \
    CP_DATA_ROOT="$PWD/in2_data/arm_$ARM" \
    mpirun --allow-run-as-root -np $NP ./build/bluscriptCP_exec > "$L" 2>&1 \
    || die "arm $ARM training failed -- see $L"
  RN=$(grep -oE 'CP_Training_log[0-9]+' "$L" | head -1 | grep -oE '[0-9]+')
  [ -n "$RN" ] || die "arm $ARM: could not determine run number"
  # A zero exit code is NOT proof the checkpoint exists: bluscriptCP catches a failed
  # save, prints [WARN], and trains on. Verify the artifact, not the signal.
  CKF=$(ls -t checkpoints_bluscriptcp/rank_0/blumodelcp_run${RN}_step_*.ckpt 2>/dev/null | head -1)
  [ -n "$CKF" ] || die "arm $ARM (run $RN) finished but wrote NO checkpoint -- check $L for 'checkpoint save ... failed'"
  say "  arm $ARM checkpoint: $(basename "$CKF") ($(du -h "$CKF" | cut -f1))"
  if grep -q 'checkpoint save.*failed' "$L"; then
    die "arm $ARM: a checkpoint save FAILED during training (see $L). Refusing to continue on a partial arm."
  fi
  echo "$RN" > "$MK"; RUNNO[$ARM]=$RN
  say "  arm $ARM done -> run $RN, final loss $(grep -E '^step ' "$L" | tail -1 | grep -oE 'loss: [0-9.]+')"
done

# ---- 3b. LEARNABILITY GATE (arm C) --------------------------------------
# The plan is explicit: if arm C's synthetic val_loss does not fall, the run says
# nothing about IN2 and must NOT be written up as a null.
CSV=CP_BluScreipt_Training_logs/CP_Training_log${RUNNO[c]}.csv
if [ -f "$CSV" ]; then
  # val_loss is forward-filled and val_freq defaults to 250, so a 400-step arm yields
  # as few as ONE real evaluation. Comparing first to last of a forward-filled column
  # then reports "did not fall" for a run that was never measured twice -- which is
  # exactly what aborted the first attempt AFTER all three arms had trained.
  NV=$(awk -F, 'NR>1 && $3+0>0 {print $3}' "$CSV" | sort -u | wc -l)
  V=$(awk -F, 'NR>1 && $3+0>0 && $3!=p {v=$3; if(f==""){f=$3} p=$3} END{print f" "v}' "$CSV")
  VF=${V% *}; VL=${V#* }
  say "arm C val_loss: $NV distinct point(s), first=$VF last=$VL"
  if [ "${NV:-0}" -lt 2 ]; then
    say "  learnability gate SKIPPED: only $NV validation point(s) -- cannot judge a trend."
    say "  (lower CP_VAL_FREQ, or lengthen the arm, to make this gate meaningful)"
    VF=1; VL=0   # force the comparison below to pass
  fi
  $PY -c "import sys; sys.exit(0 if float('$VL') < float('$VF') else 1)" \
    || die "LEARNABILITY GATE FAILED: arm C val_loss did not fall ($VF -> $VL). The run says nothing about IN2."
  say "  learnability gate: PASS"
else say "  WARN: no CSV for arm C run ${RUNNO[c]}; learnability gate skipped"; fi

# ---- 4. probe each arm (same YaRN cache it trained under) ---------------
for ARM in a b c; do
  OUTD=probe_arm_${ARM}
  [ -f $OUTD/reach.csv ] && { say "probe $ARM already done"; continue; }
  CK=$(ls -t checkpoints_bluscriptcp/rank_0/blumodelcp_run${RUNNO[$ARM]}_step_*.ckpt 2>/dev/null | head -1)
  [ -n "$CK" ] || die "arm $ARM: no checkpoint for run ${RUNNO[$ARM]}"
  say "probing arm $ARM  ($CK)"
  CKPT="$PWD/$CK" OUT=$OUTD FACTORS=$YARN CACHES=yarn YARN_SCALE=$YARN GPUS=$GPUS BLOCKS=3 \
    ./probe_1b.sh >>"$MASTER" 2>&1 || say "  WARN: probe $ARM returned nonzero"
  [ -f $OUTD/reach.csv ] && say "  arm $ARM reach.csv written" || say "  WARN: no reach.csv for arm $ARM"
done

# ---- 5. readout ---------------------------------------------------------
say "readout"
$PY Tests/bluscriptcp/in2_report.py --needle-dir probe_arm_c/needles --T $T \
    --arms a b c --R $R_MEASURED --blocks 3 > $LOGDIR/in2_report.txt 2>&1 \
    || say "  WARN: in2_report returned nonzero (see $LOGDIR/in2_report.txt)"
tail -30 $LOGDIR/in2_report.txt | tee -a "$MASTER"

# ---- 6. UPLOAD everything needed to carry forward -----------------------
say "staging artefacts for HF"
cp -r $LOGDIR $STAGE/ 2>/dev/null
for d in probe_run8 probe_run8_s4 probe_run7_yarn4 probe_arm_a probe_arm_b probe_arm_c; do
  [ -d "$d" ] && cp -r "$d" $STAGE/ 2>/dev/null
done
mkdir -p $STAGE/train_csv && cp CP_BluScreipt_Training_logs/*.csv $STAGE/train_csv/ 2>/dev/null
mkdir -p $STAGE/run_logs  && cp run_logs/*.log $STAGE/run_logs/ 2>/dev/null
cp -r /root/claude-vault $STAGE/vault 2>/dev/null
# code: git holds it, but a tarball costs nothing and survives a lost branch
tar czf $STAGE/cp_source.tgz --exclude='build' --exclude='checkpoints_bluscriptcp' \
    --exclude='in2_data' --exclude='Data_Loader/Data' --exclude='.git' --exclude='*.bin' \
    -C "$ROOT" dist/Context_Parallelism dist/zero dist/communication/src/processGroupNccl.cpp 2>/dev/null
# the three arm checkpoints (weights only -- flat, loadable at any np)
mkdir -p $STAGE/checkpoints
for ARM in a b c; do
  CK=$(ls -t checkpoints_bluscriptcp/rank_0/blumodelcp_run${RUNNO[$ARM]:-0}_step_*.ckpt 2>/dev/null | head -1)
  [ -n "$CK" ] && cp "$CK" $STAGE/checkpoints/arm_${ARM}_$(basename "$CK")
done
[ -f "$BASE_CKPT" ] && cp "$BASE_CKPT" $STAGE/checkpoints/extended_$(basename "$BASE_CKPT")
du -sh $STAGE/* 2>/dev/null | tee -a "$MASTER"

if [ -x "$HF" ]; then
  say "uploading to $HF_REPO"
  "$HF" upload "$HF_REPO" $STAGE/cp_source.tgz  source/cp_source_in2.tgz      >>"$MASTER" 2>&1
  "$HF" upload "$HF_REPO" $STAGE/vault          docs/                          >>"$MASTER" 2>&1
  "$HF" upload "$HF_REPO" $STAGE/in2_logs       logs/in2/                      >>"$MASTER" 2>&1
  "$HF" upload "$HF_REPO" $STAGE/train_csv      logs/in2_train_csv/            >>"$MASTER" 2>&1
  "$HF" upload "$HF_REPO" $STAGE/run_logs       logs/in2_stdout/               >>"$MASTER" 2>&1
  for d in probe_run8 probe_run8_s4 probe_run7_yarn4 probe_arm_a probe_arm_b probe_arm_c; do
    [ -d "$STAGE/$d" ] && "$HF" upload "$HF_REPO" "$STAGE/$d" "logs/$d/"       >>"$MASTER" 2>&1
  done
  "$HF" upload "$HF_REPO" $STAGE/checkpoints    checkpoints/in2/               >>"$MASTER" 2>&1
  say "upload complete"
else say "WARN: hf CLI not found at $HF -- artefacts staged in $STAGE but NOT uploaded"; fi

say "DONE. staged copy: $STAGE   master log: $MASTER"