#!/usr/bin/env bash
# =============================================================================
# Corrected finisher. The driver (mix64k_h200.sh) cannot run the evals:
# run_seg 695 waits for a checkpoint at step 695, but CP_CKPT_FREQ=100 means
# none can exist (695 % 100 = 95). So its watcher never fired, training ran
# straight to CP_MAX_STEPS=1300, and the driver will die at
# "step 695 checkpoint missing". The training itself is CORRECT and complete --
# only the eval orchestration is lost.
#
# This runs the evals itself off the janitor-preserved keep_step_*.ckpt, which
# exist at every multiple of 100 precisely because the janitor copies rank_0
# before CP_CKPT_KEEP can prune it. Eval loads rank_0 only (NP=1), so that is
# all it needs.
#
#   1. wait for the trainer to exit (eval needs the GPUs)
#   2. eval step 1200 at 4 lengths, one GPU each
#   3. plot+upload as 32,744 lands, again when 65,512 lands
#   4. cliff verdict; if CLIFF, also eval the last pre-cliff checkpoint
#   5. upload everything byte-verified, print a teardown verdict
# =============================================================================
set -uo pipefail
cd "$(dirname "$0")"
source /root/Adhi/BluTrain/envv/bin/activate
export LD_LIBRARY_PATH="/root/Adhi/BluTrain/Tensor-Implementations/lib:/root/Adhi/BluTrain/Profiler/lib:${LD_LIBRARY_PATH:-}"
PY=python; HF=hf
REPO=unparallelled/BluScriptCP
LOG=mix64k_logs; CK=checkpoints_bluscriptcp/rank_0
T=65536; ORIG=32768; SCALE=2
DATA=Data_Loader/Data/mix50
HAY="$PWD/Data_Loader/Data/b2_8192_16384/val_shard_00000.bin"
F="$LOG/finish.log"; : > "$F"
say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$F" >&2; }

up(){ local src=$1 dst=$2
  [ -e "$src" ] || { say "  MISSING (not uploaded): $src"; return 1; }
  $HF upload "$REPO" "$src" "$dst" >>"$LOG/upload.log" 2>&1 || { say "  UPLOAD FAILED $dst"; return 1; }
  if [ -f "$src" ]; then
    $PY - "$REPO" "$src" "$dst" <<'PYEOF' | tee -a "$F"
import sys, os
from huggingface_hub import HfApi
repo, loc, rem = sys.argv[1], sys.argv[2], sys.argv[3]
want=os.path.getsize(loc)
got={f.rfilename:(f.size or 0) for f in HfApi().repo_info(repo,files_metadata=True).siblings}.get(rem)
print(("  VERIFIED " if got==want else "  *** MISMATCH/ABSENT ")+f"{rem} local={want:,} remote={got}")
sys.exit(0 if got==want else 1)
PYEOF
  else say "  sent (dir) $dst"; fi
}

plot(){ $PY mk_mix64k_plots.py --tag "$1" --out mix64k_results >>"$F" 2>&1 \
        && say "  plotted mix64k_results_$1.{png,html}" || say "  PLOT FAILED ($1)"; }

run_eval(){ # $1 = tag   $2 = ckpt basename in rank_0
  local TAG=$1 CKF=$2
  [ -f "$CK/$CKF" ] || { say "eval $TAG: $CKF absent, SKIP"; return 1; }
  say "EVAL $TAG from $CKF -- 4 lengths, one GPU each"
  for LEN in 8192 16360 32744 65512; do
    ( GEN_DATA_ROOT=$DATA HAY="$HAY" PY=$PY OUT=mix64k_${TAG}_len_$LEN TASKS="niah_single" \
      CKPTS="m$TAG" N=100 PROMPT_LEN=$LEN MAX_NEW=24 T=$T \
      DEPTHS="0.0,0.25,0.5,0.75,1.0" YARN_ORIG=$ORIG YARN=$SCALE \
      GPUS=$(( (LEN==8192)?0:(LEN==16360)?1:(LEN==32744)?2:3 )) JOBS_PER_GPU=1 \
      BF16=1 GREEDY=1 CKPT_EXTRA="m$TAG=$CKF" ./run_longeval.sh \
      >> "$LOG/eval_${TAG}_$LEN.log" 2>&1 ) &
  done
  # plot as soon as 32,744 scores, without waiting for the slow 65,512 job
  ( for i in $(seq 1 400); do
      [ -s "mix64k_${TAG}_len_32744/summary.csv" ] && break; sleep 20; done
    for L in 8192 16360 32744; do
      v=$(sed -n 2p "mix64k_${TAG}_len_$L/summary.csv" 2>/dev/null); say "  $TAG len $L -> ${v:-pending}"; done
    plot "$TAG"
    up "mix64k_results_${TAG}.png"  "plots/mix64k_results_${TAG}.png"
    up "mix64k_results_${TAG}.html" "plots/mix64k_results_${TAG}.html"
    say "  EARLY PLOT for $TAG is on HF (8k/16k/32k)" ) &
  wait
  for LEN in 8192 16360 32744 65512; do
    v=$(sed -n 2p "mix64k_${TAG}_len_$LEN/summary.csv" 2>/dev/null)
    say "  FINAL $TAG len $LEN -> ${v:-NO RESULT}"
  done
  plot "$TAG"
  up "mix64k_results_${TAG}.png"  "plots/mix64k_results_${TAG}.png"
  up "mix64k_results_${TAG}.html" "plots/mix64k_results_${TAG}.html"
}

say "=== corrected finisher armed ==="
say "waiting for the trainer to exit (it runs to step 1300, then frees the GPUs)"
for i in $(seq 1 900); do
  pgrep -f 'build/bluscriptCP_exe[c]' >/dev/null || break
  sleep 20
done
sleep 20
say "trainer exited; last step: $(grep -aE '^step +[0-9]+' "$LOG/train.log" | tail -1 | grep -oE '^step +[0-9]+')"
say "preserved checkpoints: $(ls $CK/keep_step_*.ckpt 2>/dev/null | wc -l)"

# ---- FINAL checkpoint first. Eval 1200 only if the final looks BAD. --------
FIN=$(ls $CK/keep_step_*.ckpt 2>/dev/null | sed -E 's/.*keep_step_([0-9]+)\.ckpt/\1/' | sort -n | tail -1)
[ -n "$FIN" ] || { say "no preserved checkpoint to evaluate"; FIN=""; }
if [ -n "$FIN" ]; then
  say "=== evaluating the FINAL checkpoint, step $FIN ==="
  run_eval "$FIN" "keep_step_${FIN}.ckpt"
  say "--- verdict for step $FIN against the parent (67/65/52/29) ---"
  $PY /root/verdict.py "$FIN" 2>&1 | tee -a "$F"
  VER=$($PY /root/verdict.py "$FIN" 2>/dev/null | grep -oE '^VERDICT [A-Z]+' | awk '{print $2}')
  say "verdict: ${VER:-UNKNOWN}"
  if [ "$VER" = "BAD" ] || [ "$VER" = "UNKNOWN" ]; then
    say "=== final is ${VER} -- evaluating step 1200 to bracket where it turned ==="
    [ -f "$CK/keep_step_1200.ckpt" ] && run_eval 1200 keep_step_1200.ckpt \
      || say "  keep_step_1200.ckpt absent"
  else
    say "=== final is GOOD -- skipping the 1200 eval (saves ~38 min of GPU) ==="
  fi
fi

read -r V O B C <<<"$($PY /root/cliff.py "$LOG/train.log")"
say "cliff detector on the MIX run: $V (onset=$O baseline=$B current=$C)"
say "  for contrast, 100% b5 from the SAME parent: CLIFF at step 1199, ended 45/17/2/1"
if [ "$V" = "CLIFF" ]; then
  PRE=$(( ((O - 49) / 100) * 100 ))
  say "cliff at $O -- also evaluating the pre-cliff checkpoint at $PRE"
  [ -f "$CK/keep_step_${PRE}.ckpt" ] && run_eval "pre${PRE}" "keep_step_${PRE}.ckpt"
fi

say "=== uploading everything ==="
for d in mix64k_*_len_*; do [ -d "$d" ] && up "$d" "logs/$d/"; done
up "$LOG/master.log" "logs/mix64k/master.log"
up "$LOG/train.log"  "logs/mix64k/train.log"
up "$LOG/probe.log"  "logs/mix64k/probe.log"
up "$F"              "logs/mix64k/finish.log"
for f in mix64k_h200.sh finish2_mix64k.sh mk_mix64k_plots.py mix_buckets.py; do up "$f" "scripts/$f"; done
up "$DATA/MANIFEST.txt" "docs/mix50_MANIFEST.txt"
say "=== uploading model checkpoints (rank_0; eval and CP_INIT_FROM need only this) ==="
for k in $CK/keep_step_*.ckpt; do
  [ -e "$k" ] || continue
  s=$(basename "$k" | grep -oE '[0-9]+')
  up "$k" "checkpoints/mix64k_step_${s}.ckpt"
done
BAD=$(grep -c 'MISMATCH/ABSENT\|UPLOAD FAILED\|MISSING (not uploaded)' "$F" 2>/dev/null || echo 0)
[ "$BAD" -eq 0 ] && say "=== ALL UPLOADS VERIFIED -- safe to tear down ===" \
                 || say "=== $BAD UPLOAD PROBLEM(S) -- DO NOT TEAR DOWN ==="