#!/usr/bin/env bash
# =============================================================================
# eval_run7c.sh -- evaluate the CORRECTED-CACHE models generatively and produce
# the data for a direct comparison against the 2026-09-02 buggy-cache results
# recorded in EVAL_COMPARISON.md.
#
# THE COMPARISON THIS FEEDS
#   old run7  + buggy cache     (matched)    -> EVAL_COMPARISON.md sections 1,2,4
#   run7'     + corrected cache (matched)    -> measured here
# Both arms of that comparison are train/eval MATCHED, which is the only way the
# cache fix can be credited or blamed for a difference. The intermediate
# condition (old weights + corrected cache) is a MISMATCH and is not the test.
#
# TASKS: niah_single reproduces EVAL_COMPARISON.md section 1 (and, via DEPTHS,
# section 4); val_backward reproduces section 2. Restricting to these two keeps
# the sweep to ~2000 prompts instead of ~6000 -- the other four tasks were all at
# the floor at 16k last time and were never run in the length sweep, so adding
# them would cost 3x for cells that have no old counterpart to compare against.
#
# WAITS for arm B rather than killing it: arm B's run-number collision with the
# old run7 destroyed its step-200 shard on rank_0, so it cannot be resumed. It is
# cheaper to let it finish (~12 min) than to restart it from step 0 (~41 min),
# and finishing means arm B can be evaluated in the same sweep for free.
# =============================================================================
set -uo pipefail
cd "$(dirname "$0")"

L=eval_logs; mkdir -p "$L"
M="$L/eval_master.log"
PY=${PY:-/root/Adhi/BluTrain/envv/bin/python}
CK=checkpoints_bluscriptcp/rank_0
LENS=${LENS:-"1024 2048 4096 8192 16360"}
N=${N:-100}
TASKS=${TASKS:-"niah_single val_backward"}

# say() writes to STDERR. It used to write to stdout, which silently poisoned
# every $(...) capture of a function that also logged -- that is what fed a
# 3-line checkpoint path into arm B's CP_INIT_FROM and aborted it on rank 2.
say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$M" >&2; }
die(){ say "FATAL: $*"; exit 1; }

BASE_CK=blumodelcp_run6_step_750.ckpt        # run7', the corrected-cache 16k base
ARMB_RUN=$(cat retrain_logs/arm_b.runnum 2>/dev/null || echo 7)

# ---- 1. wait for the GPUs -------------------------------------------------
if pgrep -f bluscriptCP_exec >/dev/null 2>&1; then
  say "arm B still training -- waiting for it to finish (not killing it: its"
  say "  step-200 shard on rank_0 was lost to a run-number collision, so a"
  say "  resume is impossible and a kill would cost 41 min, not 12)"
  for i in $(seq 1 240); do
    pgrep -f bluscriptCP_exec >/dev/null 2>&1 || break
    sleep 15
  done
fi
pgrep -f bluscriptCP_exec >/dev/null 2>&1 && die "training still running after 60 min"
say "GPUs free"

# ---- 2. register whatever checkpoints actually exist ------------------------
EXTRA=""
[ -f "$CK/$BASE_CK" ] || die "run7' checkpoint missing: $CK/$BASE_CK"
EXTRA="base_c=$BASE_CK"
say "registered base_c = $BASE_CK  (run7', corrected cache)"

ARMB_CK=$(ls -t $CK/blumodelcp_run${ARMB_RUN}_step_*.ckpt 2>/dev/null | head -1)
if [ -n "$ARMB_CK" ]; then
  # all four ZeRO-1 shards must be present or the load is silently partial
  n=$(ls checkpoints_bluscriptcp/rank_*/"$(basename "$ARMB_CK")" 2>/dev/null | wc -l)
  if [ "$n" -ge 4 ]; then
    EXTRA="$EXTRA armB_c=$(basename "$ARMB_CK")"
    say "registered armB_c = $(basename "$ARMB_CK")  (all 4 shards present)"
  else
    say "SKIP armB_c: only $n/4 rank shards for $(basename "$ARMB_CK")"
  fi
else
  say "SKIP armB_c: no run${ARMB_RUN} checkpoint"
fi

MODELS=$(echo $EXTRA | tr ' ' '\n' | cut -d= -f1 | tr '\n' ' ')
say "models: $MODELS"

# ---- 3. the sweep ----------------------------------------------------------
# YARN_ORIG=4096 is MATCHED here: these weights were trained under it.
for LEN in $LENS; do
  say "eval PROMPT_LEN=$LEN  N=$N  tasks: $TASKS"
  PY=$PY OUT=newck_len_$LEN TASKS="$TASKS" CKPTS="$MODELS" N=$N \
  PROMPT_LEN=$LEN MAX_NEW=24 T=16384 DEPTHS="0.0,0.25,0.5,0.75,1.0" \
  YARN_ORIG=4096 GPUS=0,1,2,3 JOBS_PER_GPU=2 BF16=1 GREEDY=1 \
  CKPT_EXTRA="$EXTRA" ./run_longeval.sh >> "$L/eval_$LEN.log" 2>&1 \
    || say "  WARN: eval at $LEN returned nonzero (see $L/eval_$LEN.log)"
  # report as we go, so a partial sweep is still usable
  if [ -f "newck_len_$LEN/summary.csv" ]; then
    say "  -> $(grep -c . newck_len_$LEN/summary.csv) rows in newck_len_$LEN/summary.csv"
  else
    say "  -> NO summary.csv at $LEN"
  fi
done

say "EVAL DONE"
for LEN in $LENS; do
  [ -f "newck_len_$LEN/summary.csv" ] && say "  newck_len_$LEN/summary.csv OK" \
                                      || say "  newck_len_$LEN MISSING"
done