#!/usr/bin/env bash
# Phase-1 eval, retry. Two fixes over the first attempt:
#   1. GEN_DATA_ROOT -- run_longeval.sh defaults CP_DATA_ROOT to Data_Loader/Data,
#      which on this box holds only SUBDIRECTORIES. The trainer builds a data
#      loader even for pure generation, so it aborted with "no .bin shards found
#      for split train" before generating a single token.
#   2. logs/upload.log is uploaded as a SNAPSHOT copy -- up() appends to it while
#      uploading it, so the live file grows between upload and verification and
#      always reports a size mismatch.
set -uo pipefail
cd "$(dirname "$0")"
export PATH=/root/Adhi/BluTrain/envv/bin:$PATH
source /root/Adhi/BluTrain/envv/bin/activate 2>/dev/null
REPO=unparallelled/BluScriptCP
L=p1_logs; M="$L/eval2.log"
say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$M" >&2; }
RN=$(cat "$L/runnum"); FSTEP=1200
FINAL="blumodelcp_run${RN}_step_${FSTEP}.ckpt"
[ -f "checkpoints_bluscriptcp/rank_0/$FINAL" ] || { say "FATAL: $FINAL missing"; exit 1; }
A1B_680="CP_N_EMBD=1536 CP_N_LAYER=24 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=1"

for PL in 8192 16360; do
  O="p1_eval_len${PL}"
  say "eval prompt_len=$PL  (p1final @1200 + base680 zero-shot-YaRN baseline)"
  A1B="$A1B_680" HAY="Data_Loader/Data/b2_8192_16384/val_shard_00000.bin" \
  GEN_DATA_ROOT="Data_Loader/Data/p1mix" \
  CKPT_EXTRA="p1final=$FINAL base680=varybatch.ckpt" \
  CKPTS="p1final base680" TASKS=niah_single N=100 \
  T=16384 YARN=4 YARN_ORIG=4096 PROMPT_LEN=$PL GPUS=0,1 NP=1 OUT="$O" \
    ./run_longeval.sh >> "$L/eval2_${PL}.log" 2>&1
  if [ -s "$O/summary.csv" ] && [ "$(wc -l < "$O/summary.csv")" -gt 1 ]; then
    say "  len $PL -> $(tail -1 "$O/summary.csv")"
  else
    say "  len $PL FAILED"; grep -aA2 'what()' "$O"/gen/*.log 2>/dev/null | head -4 >&2
  fi
done

say "=== PHASE 1 NIAH RESULTS ==="
for PL in 8192 16360; do
  f="p1_eval_len${PL}/summary.csv"
  [ -f "$f" ] && { say "  len $PL: header=$(head -1 $f)  values=$(tail -1 $f)"; }
done

say "regenerating plots with eval points"
python mk_p1_plots.py 2>&1 | tee -a "$M"

up(){ [ -f "$1" ] || { say "  MISSING $1"; return 1; }
  hf upload "$REPO" "$1" "$2" >>"$L/upload2.log" 2>&1
  python - "$REPO" "$1" "$2" <<'VEOF'
import sys,os
from huggingface_hub import HfApi
w=os.path.getsize(sys.argv[2])
g={f.rfilename:(f.size or 0) for f in HfApi().repo_info(sys.argv[1],files_metadata=True).siblings}.get(sys.argv[3])
ok=(g==w); print(("  VERIFIED  " if ok else "  *** FAILED ")+f"{sys.argv[3]} local={w:,} remote={g}")
sys.exit(0 if ok else 1)
VEOF
}
BAD=0
say "=== uploading eval + plots ==="
for PL in 8192 16360; do
  for f in p1_eval_len${PL}/summary.csv p1_eval_len${PL}/depth_summary.csv \
           p1_eval_len${PL}/scores/p1final_niah_single.csv \
           p1_eval_len${PL}/scores/base680_niah_single.csv \
           p1_eval_len${PL}/gen/p1final_niah_single.csv; do
    [ -f "$f" ] && { up "$f" "evals/p1_len${PL}_$(echo $f | sed 's|.*/||')" || BAD=$((BAD+1)); }
  done
  [ -f "$L/eval2_${PL}.log" ] && { up "$L/eval2_${PL}.log" "logs/p1_eval2_${PL}.log" || BAD=$((BAD+1)); }
done
up plots_p1/p1_lineage.png plots/p1_680m_lineage.png || BAD=$((BAD+1))
up plots_p1/p1_results.png plots/p1_680m_results.png || BAD=$((BAD+1))
up mk_p1_plots.py scripts/mk_p1_plots.py || BAD=$((BAD+1))
up p1_eval2.sh scripts/p1_eval2.sh || BAD=$((BAD+1))
cp "$L/upload.log" "$L/upload_snapshot.log" 2>/dev/null
up "$L/upload_snapshot.log" "logs/p1_upload.log" || BAD=$((BAD+1))
up "$M" "logs/p1_eval2.log" || BAD=$((BAD+1))

say "=== FINAL VERDICT ==="
GOT=0; for PL in 8192 16360; do [ -s "p1_eval_len${PL}/summary.csv" ] && \
  [ "$(wc -l < p1_eval_len${PL}/summary.csv)" -gt 1 ] && GOT=$((GOT+1)); done
if [ "$BAD" -eq 0 ] && [ "$GOT" -eq 2 ]; then
  say "training + checkpoint + evals ($GOT/2) + plots ALL VERIFIED -- SAFE TO TEAR DOWN"
else
  say "*** $BAD upload failure(s), $GOT/2 evals -- DO NOT TEAR DOWN ***"
fi