#!/usr/bin/env bash
# =============================================================================
# THE CONTROL: is the ceiling a property of the CONFIG or of the 680M?
#
# mix64k_step1200 (1.06B, best of the 3-stage ladder) was TRAINED at
# ORIG=32768 SCALE=2, which puts a 16,360 prompt at virtual 8,180 -- twice
# beyond the 4,096 range it was pretrained on. Re-parameterising to
# ORIG=4096 SCALE=16 (same T=65,536) puts it at virtual 1,022.
#
# mixcont_g1400_cliff is the COLLAPSED checkpoint. If the collapse is partly a
# config artefact it should recover too; if it is real damage it should not.
#
# PRE-REGISTERED PREDICTIONS (mix64k), written before the run:
#   16,360 @ x16 -> 85-92   (known 70 at its trained config)
#   32,744 @ x32 -> 70-80   (known 50)
#   65,512 @ x32 -> 40-50   (known 31)
#   16,360 @ x2/ORIG=32768 -> 66-74  <- REPRODUCTION CHECK of the prior 70.
#      If this does not reproduce, the harness differs from the prior session
#      and NONE of the cross-session comparisons in this analysis are valid.
# =============================================================================
set -uo pipefail
cd "$(dirname "$0")"
export PATH=/root/Adhi/BluTrain/envv/bin:$PATH
source /root/Adhi/BluTrain/envv/bin/activate 2>/dev/null
L=ctl_logs; mkdir -p "$L"; M="$L/master.log"
say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$M" >&2; }
CK=checkpoints_bluscriptcp/rank_0
HAY=Data_Loader/Data/b2_8192_16384/val_shard_00000.bin
[ "$(stat -c%s "$HAY")" = 199754694 ] || { say "FATAL haystack differs"; exit 1; }

for f in mix64k_step_1200 mixcont_g1400_cliff; do
  [ -f "$CK/$f.ckpt" ] && continue
  say "downloading $f.ckpt"
  hf download unparallelled/BluScriptCP --quiet --include "checkpoints/$f.ckpt" --local-dir dl_c || exit 1
  mv -f "dl_c/checkpoints/$f.ckpt" "$CK/" || exit 1
done
rm -rf dl_c
for f in mix64k_step_1200 mixcont_g1400_cliff; do
  [ "$(stat -c%s "$CK/$f.ckpt")" = 3750897182 ] || { say "FATAL $f wrong size"; exit 1; }
done
say "both 1.06B checkpoints verified at 3,750,897,182 B"

# 1.06B arch is run_longeval.sh's DEFAULT A1B (36 layers, weight_tying=0) -- do not override
run(){ # T SCALE ORIG PL tag
  local T=$1 S=$2 OR=$3 PL=$4 TAG=$5
  local O="ctl_${TAG}_len${PL}"
  [ $((S*OR)) -eq $T ] || { say "YaRN invariant $S*$OR != $T"; return 1; }
  [ -s "$O/summary.csv" ] && [ "$(wc -l < "$O/summary.csv")" -gt 1 ] && {
    say "  $TAG len $PL cached: $(grep '^niah_single,' $O/summary.csv)"; return 0; }
  say "  $TAG len $PL  (T=$T scale=$S orig=$OR -> virtual $((PL/S)))"
  HAY="$HAY" GEN_DATA_ROOT="Data_Loader/Data/evalsrc" \
  CKPT_EXTRA="mix1200=mix64k_step_1200.ckpt cliff1400=mixcont_g1400_cliff.ckpt" \
  CKPTS="mix1200 cliff1400" TASKS=niah_single N=100 \
  T=$T YARN=$S YARN_ORIG=$OR PROMPT_LEN=$PL GPUS=0,1 NP=1 OUT="$O" \
    ./run_longeval.sh >> "$L/${TAG}_${PL}.log" 2>&1
  if [ -s "$O/summary.csv" ] && [ "$(wc -l < "$O/summary.csv")" -gt 1 ]; then
    say "    -> $(head -1 $O/summary.csv) | $(grep '^niah_single,' $O/summary.csv)"
  else say "    FAILED"; grep -ahA2 'what()' "$O"/gen/*.log 2>/dev/null|head -3|tee -a "$M" >&2; fi
}
say "=== 1.06B CONTROL: config vs model ==="
say "REPRODUCTION CHECK first -- if 16,360 @ x2/orig32768 is not ~70, stop trusting everything"
run 65536  2 32768 16360 astrained
say "now the corrected parameterisation (ORIG held at the 4096 pretrain length)"
run 65536 16  4096 16360 x16
run 131072 32 4096 32744 x32
run 131072 32 4096 65512 x32
say "=== CONTROL RESULTS (cols: mix1200, cliff1400) ==="
for t in astrained_16360 x16_16360 x32_32744 x32_65512; do
  f="ctl_${t%_*}_len${t##*_}/summary.csv"
  [ -f "$f" ] && say "  $t  $(head -1 $f) -> $(grep '^niah_single,' $f)"
done
say "DONE"