#!/bin/bash
# Nex-N2.5-mini phase 2b. Phase 2 stopped at 22:24Z: this build's ROCm0 path computes the BF16 MoE GGUF wrong
# (wikitext PPL 106 vs 7.8 for its own 4-bit quant; -fa off and "-ngl 0" with the default op offload are wrong too).
# Controls, chunk 1: pure CPU (-dev none --no-op-offload) 5.6964, Vulkan0 5.6953.
# => BF16 reference logits and the imatrix are computed on PURE CPU; every 4-bit file is graded on BOTH GPU backends
#    against that reference; then speed + gates -> memory sizing -> Nex FAST seats -> un-ice (as phase 2).
set -uo pipefail
W=/mnt/models/nex-n2.5-mini; A=/mnt/models/agnes-3.0-flash; B=/opt/llama-rocm/rocmfpx-724/build-hipvk/bin; N=Nex-N2.5-mini
export LD_LIBRARY_PATH=$B:/opt/rocm-7.2.4/lib HSA_OVERRIDE_GFX_VERSION=11.5.1 GGML_HIP_ENABLE_UNIFIED_MEMORY=1
cd $W; mkdir -p kld imat out-imat results
prlimit --pid $$ --core=1:1   # RLIMIT_CORE of exactly 1 BYTE (bash `ulimit -c 1` is 1 KiB and does not stop apport)
log(){ echo "[$(date -u +%FT%TZ)] $*"; }
table(){
  python3 - "$1" <<'EOF'
import sys
out = []
for l in open(sys.argv[1], errors="replace"):
    t = l.split()
    if len(t) > 2 and t[0].isdigit() and l.rstrip().endswith("%") and l.count("%") == 2:
        out.append(" ".join(t))
print("\n".join(out))
EOF
}
chunk1(){ grep -oE '\[1\][0-9.]+' "$1" | head -1 | cut -d']' -f2; }
CPU="-dev none -ngl 0 --no-op-offload -t 16"
PPL="$B/llama-perplexity -c 2048 -b 2048"
BF=gguf/$N-BF16.gguf; Q=$B/llama-quantize; TXT=$A/calib/wikitext-2-raw/wiki.test.raw

log "N1c BF16 perplexity + base logits on pure CPU (wikitext-2 test, 40 chunks)"
$PPL -m $BF -f $TXT --chunks 40 --kl-divergence-base kld/bf16.kld $CPU > logs/N1c_ppl_bf16_cpu.log 2>&1
log "N1c exit=$? $(grep -E 'Final estimate' logs/N1c_ppl_bf16_cpu.log)"
a=$(chunk1 logs/N1c_ppl_bf16_cpu.log); b=$(chunk1 logs/diag_bf16_vk_faon.log)
if ! python3 -c "import sys; a, b = float('${a:-nan}'), float('${b:-nan}'); sys.exit(0 if abs(a - b) / b < 0.005 else 1)"; then
  log "GATE FAIL: CPU chunk-1 PPL ${a:-missing} vs Vulkan0 control ${b:-missing} -> stop (box stays iced)"; log "NEX_PHASE2B_FAILED"; exit 1
fi
echo "{\"step\":\"N1c\",\"cpu_chunk1\":$a,\"vulkan0_chunk1\":$b,\"hip_rocm0_chunk1\":$(chunk1 logs/N1_ppl_bf16.log),\"hip_rocm0_faoff_chunk1\":$(chunk1 logs/diag_bf16_rocm_faoff.log),\"hip_rocm0_final\":$(grep -oE 'Final estimate: PPL = [0-9.]+' logs/N1_ppl_bf16.log | grep -oE '[0-9.]+$'),\"result\":\"PASS\"}" | tee results/nex_reference.jsonl
log "N1c gate PASS (CPU $a vs Vulkan0 $b)"

log "N2c imatrix on pure CPU (BF16, bartowski calibration_datav3, c=512)"
$B/llama-imatrix -m $BF -f $A/calib/calibration_datav3.txt -o imat/$N.imatrix -c 512 -b 512 $CPU > logs/N2c_imatrix_cpu.log 2>&1
rc=$?; log "N2c exit=$rc"
[ $rc -eq 0 ] && [ -s imat/$N.imatrix ] || { log "imatrix failed -> stop (box stays iced)"; log "NEX_PHASE2B_FAILED"; exit 2; }
sha256sum imat/$N.imatrix > imat/$N.imatrix.sha256

log "N3 imatrix tiers on CPU (background) while the standard tiers are graded on the GPU"
(
  systemd-run --scope --quiet -p MemoryMax=24G -p MemorySwapMax=0 nice -n 5 bash -c "
    $Q --imatrix imat/$N.imatrix --output-tensor-type q6_K $BF out-imat/$N-imatrix-Q4_0_ROCMFP4_STRIX_LEAN.gguf Q4_0_ROCMFP4_STRIX_LEAN 16 > logs/N3_q106i.log 2>&1
    $Q --imatrix imat/$N.imatrix --output-tensor-type q6_K --token-embedding-type q6_K $BF out-imat/$N-imatrix-Q4_0_ROCMFP4_COHERENT.gguf Q4_0_ROCMFP4_COHERENT 16 > logs/N3_q102i.log 2>&1
    $Q --imatrix imat/$N.imatrix --output-tensor-type q6_K $BF out-imat/$N-imatrix-Q4_0_ROCMFP4_FAST.gguf Q4_0_ROCMFP4_FAST 16 > logs/N3_q103i.log 2>&1"
  python3 readback.py Q6_K Q5_K out-imat/$N-imatrix-Q4_0_ROCMFP4_STRIX_LEAN.gguf >  logs/N3_readback.log
  python3 readback.py Q6_K Q6_K out-imat/$N-imatrix-Q4_0_ROCMFP4_COHERENT.gguf   >> logs/N3_readback.log
  python3 readback.py Q6_K -    out-imat/$N-imatrix-Q4_0_ROCMFP4_FAST.gguf       >> logs/N3_readback.log
  echo N3_DONE >> logs/N3_readback.log
) &
n3=$!
grade(){  # tag file device logprefix
  $PPL -m "$2" --kl-divergence-base kld/bf16.kld --kl-divergence -ngl 999 -fa on -dio -dev $3 > logs/$4_kld_$1.log 2>&1
  log "  $4 $1 exit=$? $(grep -E 'Mean +KLD' logs/$4_kld_$1.log | tr -s ' ')"
}
for t in "q106 out/$N-Q4_0_ROCMFP4_STRIX_LEAN.gguf" "q102 out/$N-Q4_0_ROCMFP4_COHERENT.gguf" "q103 out/$N-Q4_0_ROCMFP4_FAST.gguf"; do
  set -- $t; grade $1 $2 ROCm0 N4; grade $1 $2 Vulkan0 N4v
done
wait $n3; log "N3 done: $(grep -c ^PASS logs/N3_readback.log)/3 read-back PASS"; cat logs/N3_readback.log
for l in N3_q106i N3_q102i N3_q103i; do printf "%-9s " $l; grep -oE "quant size\s*=\s*[0-9.]+ MiB \([0-9.]+ BPW\)" logs/$l.log; done | tee -a logs/Q_sizes.log
log "N4 grade the imatrix tiers (both backends)"
for t in "q106i out-imat/$N-imatrix-Q4_0_ROCMFP4_STRIX_LEAN.gguf" "q102i out-imat/$N-imatrix-Q4_0_ROCMFP4_COHERENT.gguf" "q103i out-imat/$N-imatrix-Q4_0_ROCMFP4_FAST.gguf"; do
  set -- $t; grade $1 $2 ROCm0 N4; grade $1 $2 Vulkan0 N4v
done

log "N5 repeat the STRIX_LEAN grade on both backends; every per-chunk row must match"
: > results/nex_repeat.jsonl
for pair in "ROCm0 N4 N5" "Vulkan0 N4v N5v"; do
  set -- $pair
  $PPL -m out/$N-Q4_0_ROCMFP4_STRIX_LEAN.gguf --kl-divergence-base kld/bf16.kld --kl-divergence -ngl 999 -fa on -dio -dev $1 \
    > logs/$3_kld_q106_repeat.log 2>&1
  x=$(table logs/$3_kld_q106_repeat.log); y=$(table logs/$2_kld_q106.log); n=$(printf '%s\n' "$x" | grep -c .)
  echo "{\"step\":\"$3\",\"device\":\"$1\",\"rows\":$n,\"result\":\"$( [ "$n" -ge 40 ] && [ "$x" = "$y" ] && echo MATCH || echo MISMATCH )\"}" | tee -a results/nex_repeat.jsonl
done

log "N6 speed + gates"
python3 nex_bench.py > logs/N6_bench.log 2>&1; log "N6 exit=$?"
log "N7 memory sizing"
bash nex_sizing.sh > logs/N7_sizing.log 2>&1; log "N7 exit=$?"
log "N8 Nex FAST seats + un-ice"
bash $A/unice.sh > logs/N8_unice.log 2>&1; log "N8 exit=$? $(tail -1 logs/N8_unice.log)"
log "NEX_PHASE2_DONE"