#!/usr/bin/env bash
# =============================================================================
# run_longeval.sh -- batch RULER / NIAH / VAL generative evaluation for bluscriptCP.
#
# Scores what the model GENERATES (substring match), unlike probe_1b.sh which
# measures teacher-forced NLL over an answer span already present in the input.
#
# Requires a binary built with generation support:
#     make CP_FUSED_ROPE=1 CP_ATTN_SPLIT=1 SM_ARCH=90 bluscript-cp -j$(nproc)
#
# Usage:
#     ./run_longeval.sh                          # all tasks, all checkpoints
#     TASKS="niah_single vt" ./run_longeval.sh   # subset
#     CKPTS="arm_c" N=20 ./run_longeval.sh       # quick smoke
#
# Resumable: each (ckpt,task) writes a .done marker; re-running skips finished
# pairs. Kill it any time -- completed pairs survive.
# =============================================================================
set -uo pipefail
cd "$(dirname "$0")"

ROOT="$(cd ../.. && pwd)"
export LD_LIBRARY_PATH="$ROOT/Tensor-Implementations/lib:$ROOT/Profiler/lib:${LD_LIBRARY_PATH:-}"
export OMPI_ALLOW_RUN_AS_ROOT=1 OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1

EXE=./build/bluscriptCP_exec
PY=${PY:-python3}
OUT=${OUT:-longeval_runs}
CK=checkpoints_bluscriptcp/rank_0

# CP_T is the RoPE cache length. The C++ side pads each decode step up to a
# multiple of 256 (the Hopper split kernel's T_q guard), so
# pad_up(PROMPT_LEN + MAX_NEW) must be <= CP_T or it aborts up front.
T=${T:-16384}
# COST: a forward at T=16384 fp32 measures 17.8 s on H200 / ~25 s on H100
# (probe_arm_c: 21 jobs x 8 windows in 12.5 min at slots=4). Naive re-forward
# decoding costs ONE such forward per generated token, so the defaults matter:
#   fp32 MAX_NEW=32 N=50 -> 190 GPU-h    <- do not
#   bf16, early stop, 4 tasks at ~6 forwards + niah_multivalue at ~24 (its
#   MIN_TOK=14 blocks an early break) -> weighted (4*6+24)/5 = 9.6 /prompt
#   -> ~18.7 GPU-h for 1000 prompts -> ~4.7 h on 4 GPUs. Step 2 measures this
#   on BOTH niah_single and niah_multivalue rather than assuming it.
# HARD CONSTRAINT: PROMPT_LEN + MAX_NEW must be <= 16384 AND pad up to exactly it.
# bluscriptCP.cpp:2446 dies when pad_up(P_LEN+MAX_NEW) > CP_T, and PAD=256, so
# pad_up(16368+24)=16640 > 16384 would abort EVERY job. 16360+24 = 16384 exactly.
MAX_NEW=${MAX_NEW:-24}
PROMPT_LEN=${PROMPT_LEN:-16360}
N=${N:-50}
BF16=${BF16:-1}
GPUS=${GPUS:-0}          # comma list, e.g. "0,1,2,3" -> (ckpt,task) jobs fan out
NP=${NP:-1}              # CP_SIZE per job; GPUs are used for THROUGHPUT, not size
YARN=${YARN:-4}
# YARN_ORIG_MAXPOS (YARNOps.cpp:57) defaults to 1024, but the base model was
# trained at context_length 4096 -- hence the warning EVERY run emits:
#   seq_len(16384) != scale(4.000)*orig_maxpos(1024.0)=4096.0
# i.e. the YaRN ramp targets 4096 while the cache is SIZED 16384, so positions
# past ~4k sit outside the range it was parameterised for. YARN_ORIG=4096 makes
# scale*orig == seq_len exactly. CAUTION: this is a train/eval MISMATCH -- the
# checkpoints were trained under the 1024 default, so a correctly-parameterised
# cache is not one they ever saw. Diagnostic, not a guaranteed fix.
YARN_ORIG=${YARN_ORIG:-}
GREEDY=${GREEDY:-1}      # deterministic: benchmarks must be reproducible

# Needle depth grid. MUST match needle_gen.py's p_frac points or the generative
# curve cannot be overlaid on the NLL curve. Both place the needle at a fraction
# of ABSOLUTE POSITION (needle_gen p0.25 -> slot 4087 -> fraction 0.25), so the
# grids correspond directly.
DEPTHS=${DEPTHS:-"0.0,0.25,0.5,0.75,1.0"}

# Minimum substantive tokens before a newline may end generation. 3 suits a
# 7-digit number or a single word, but niah_multivalue needs THREE values
# (~12-15 tokens): at 3 a model that emits value 1 then a newline would stop with
# 1 of 3 and score MISS -- re-creating the bug the guard exists to prevent,
# precisely for the long-answer task.
declare -A MINTOK=( [niah_multivalue]=14 )

TASKS=${TASKS:-"niah_single niah_multikey niah_multivalue vt val_forward val_backward"}
# name:checkpoint-file
CKPTS=${CKPTS:-"base arm_a arm_b arm_c"}
declare -A CKPT_FILE=(
  [base]=blumodelcp_run7_step_750.ckpt
  [arm_a]=blumodelcp_run12_step_400.ckpt
  [arm_b]=blumodelcp_run15_step_400.ckpt
  [arm_c]=blumodelcp_run16_step_400.ckpt
)

A1B="CP_N_EMBD=1536 CP_N_LAYER=36 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=0"
ARCH="$A1B CP_ATTN_MODE=ring CP_ATTN_FUSION=split"

mkdir -p "$OUT"/{prompts,gen,state,scores}
M="$OUT/master.log"
say() { echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$M"; }

[ -x "$EXE" ] || { echo "ERROR: $EXE missing -- build first:"; \
  echo "  make CP_FUSED_ROPE=1 CP_ATTN_SPLIT=1 SM_ARCH=90 bluscript-cp -j\$(nproc)"; exit 1; }

# Fail early on a binary built without generation support rather than after an
# hour of prompt generation.
if ! grep -q CP_GENERATE "$EXE" 2>/dev/null; then
  echo "ERROR: $EXE has no CP_GENERATE support (stale build). Rebuild:"
  echo "  make CP_FUSED_ROPE=1 CP_ATTN_SPLIT=1 SM_ARCH=90 bluscript-cp -j\$(nproc)"
  exit 1
fi

say "longeval: T=$T prompt_len=$PROMPT_LEN max_new=$MAX_NEW n=$N greedy=$GREEDY gpus=$GPUS"

# ---- 1. prompts (shared across checkpoints; identical bytes => fair compare) --
HAY=$(ls Data_Loader/Data/flux_val_*.bin 2>/dev/null | head -1)
[ -z "$HAY" ] && HAY=$(ls Data_Loader/Data/flux_train_*.bin 2>/dev/null | head -1)
say "haystack: ${HAY:-<synthetic fallback>}"
[ -z "$HAY" ] && say "  WARN: synthetic filler is a weaker distractor than real text -- it INFLATES accuracy. Record this."

# A bare "-f the .bin exists" test silently reuses prompts from a different run:
# a smoke run at N=2 then poisons a main run at N=50 (the C++ dies on "short read
# on prompt 2" AFTER flushing 2 rows, which the old [ -s ] scoring gate happily
# scored and marked .done, permanently). The signature covers everything that
# changes prompt content -- N, length, depths, seed, AND the haystack, so a
# synthetic-filler run is not reused once the flux shard arrives.
SIG="N=$N|LEN=$PROMPT_LEN|DEPTHS=$DEPTHS|SEED=0|HAY=$(basename "${HAY:-none}")|HAYSZ=$(stat -c %s "${HAY:-/nonexistent}" 2>/dev/null || echo 0)"
if [ "$(cat "$OUT/prompts/.signature" 2>/dev/null)" != "$SIG" ]; then
  say "prompt signature changed -- regenerating prompts and clearing state"
  rm -f "$OUT"/prompts/*_prompts.bin "$OUT"/prompts/*_meta.json
  # MUST clear state too: the generation skip below tests .gen, so leaving those
  # markers would skip regeneration and score OLD tokens against the NEW
  # meta.json -- attributing accuracy to depths the tokens never came from.
  rm -f "$OUT"/state/*
fi

GEN_OK=1
for TASK in $TASKS; do
  if [ ! -f "$OUT/prompts/${TASK}_prompts.bin" ]; then
    say "generating prompts: $TASK"
    $PY Tests/bluscriptcp/longeval_gen.py --task "$TASK" --prompt-len "$PROMPT_LEN" \
        --n "$N" --seed 0 --depths "$DEPTHS" ${HAY:+--haystack "$HAY"} \
        --out-dir "$OUT/prompts" \
        >>"$M" 2>&1 || { say "  FAILED to generate $TASK"; GEN_OK=0; continue; }
  fi
  # exact stride: the C++ reader slices the stream at N x PROMPT_LEN x 4 bytes
  WANT=$(( N * PROMPT_LEN * 4 ))
  GOT=$(stat -c %s "$OUT/prompts/${TASK}_prompts.bin" 2>/dev/null || echo 0)
  [ "$GOT" = "$WANT" ] || { say "  BAD SIZE $TASK: ${GOT}B want ${WANT}B"; GEN_OK=0; }
done
# Only claim the signature if every task really produced prompts, else the next
# run would match it and never retry the missing ones.
[ "$GEN_OK" = 1 ] && echo "$SIG" > "$OUT/prompts/.signature"

# ---- 2. generate (fanned out over GPUs) -------------------------------------
# Each (checkpoint, task) pair is one independent job needing NP GPUs, so the
# GPU list is used for THROUGHPUT: slots = nGPU / NP jobs run concurrently.
# 24 pairs (4 ckpts x 6 tasks) on 8 GPUs at NP=1 is 8-way parallel.
IFS=',' read -r -a GPUARR <<< "$GPUS"
NGPU=${#GPUARR[@]}
# One job needs ~24.3 GB (measured) of an 80 GB card and drives it to only ~42%
# util -- B=1 decoding of a 1.06B model cannot saturate an H100, and memory is
# nowhere near the limit. So oversubscribe: JOBS_PER_GPU>1 runs several jobs per
# card and buys ~1.8x at 2/GPU. 3/GPU is ~73 GB of 81 -- it fits but leaves
# little headroom, so it is opt-in rather than the default.
JOBS_PER_GPU=${JOBS_PER_GPU:-1}
SLOTS=$(( (NGPU / NP) * JOBS_PER_GPU )); [ "$SLOTS" -lt 1 ] && SLOTS=1
say "fan-out: ${NGPU} GPU(s), NP=$NP -> $SLOTS concurrent job(s)"

jobs=()
for NAME in $CKPTS; do
  F="${CKPT_FILE[$NAME]:-}"
  [ -n "$F" ] || { say "unknown checkpoint key: $NAME"; continue; }
  [ -f "$CK/$F" ] || { say "SKIP $NAME: $CK/$F not present"; continue; }
  for TASK in $TASKS; do
    # Skip on .gen, not .done: .done is written only by the scoring loop, which
    # runs after `wait` for ALL generation, so a teardown mid-run would otherwise
    # lose every completed pair. Safe because a signature change clears state/.
    [ -f "$OUT/state/${NAME}.${TASK}.gen" ] && { say "  $NAME/$TASK already generated"; continue; }
    [ -f "$OUT/prompts/${TASK}_prompts.bin" ] || { say "  SKIP $NAME/$TASK: no prompts"; continue; }
    jobs+=("$NAME:$TASK")
  done
done
say "${#jobs[@]} job(s) to run"

running=0; g=0
for J in "${jobs[@]:-}"; do
  [ -z "$J" ] && continue
  NAME=${J%%:*}; TASK=${J##*:}
  F="${CKPT_FILE[$NAME]}"
  PB="$OUT/prompts/${TASK}_prompts.bin"
  G="$OUT/gen/${NAME}_${TASK}.csv"
  # g cycles over SLOTS, which now EXCEEDS the GPU count when JOBS_PER_GPU>1,
  # so the device offset must wrap modulo NGPU -- otherwise slots beyond the
  # last GPU would compute an empty CUDA_VISIBLE_DEVICES and every job there
  # would silently land on device 0.
  dev_off=$(( (g * NP) % NGPU ))
  devs=$(IFS=,; echo "${GPUARR[*]:$dev_off:$NP}")
  say "  [gpu $devs] $NAME / $TASK"
  (
    env $ARCH CP_SIZE=$NP CP_T=$T CP_B=1 YARN_SCALE=$YARN \
        ${BF16:+CP_BF16=$BF16} \
          ${YARN_ORIG:+YARN_ORIG_MAXPOS=$YARN_ORIG} \
        CP_INIT_FROM="$PWD/$CK/$F" \
        CP_GENERATE=1 CP_GEN_TOKENS_BIN="$PWD/$PB" \
        CP_GEN_PROMPT_LEN=$PROMPT_LEN CP_GEN_N=$N CP_GEN_MAX_NEW=$MAX_NEW \
        CP_GEN_GREEDY=$GREEDY CP_GEN_OUT="$PWD/$G" \
        CP_GEN_MIN_TOKENS=${MINTOK[$TASK]:-3} \
        CUDA_VISIBLE_DEVICES=$devs \
        mpirun --allow-run-as-root -np $NP "$EXE" \
        > "$OUT/gen/${NAME}_${TASK}.log" 2>&1 \
      && [ -s "$G" ] && touch "$OUT/state/${NAME}.${TASK}.gen"
  ) &
  g=$(( (g + 1) % SLOTS )); running=$(( running + 1 ))
  if [ "$running" -ge "$SLOTS" ]; then wait -n; running=$(( running - 1 )); fi
done
wait

# ---- 2b. score (cheap, CPU-only, after all generation) ----------------------
for NAME in $CKPTS; do
  for TASK in $TASKS; do
    G="$OUT/gen/${NAME}_${TASK}.csv"
    MK="$OUT/state/${NAME}.${TASK}.done"
    [ -f "$MK" ] && continue
    if [ ! -s "$G" ]; then
      [ -f "$OUT/gen/${NAME}_${TASK}.log" ] && \
        say "  NO OUTPUT $NAME/$TASK -- see $OUT/gen/${NAME}_${TASK}.log"
      continue
    fi
    # A non-empty file is NOT a complete one. If the C++ died partway (e.g. a
    # short read on a stale prompt file) it still flushed the rows it managed,
    # and scoring those would report an accuracy over a handful of prompts and
    # then mark the pair .done. Note: `$(grep -c . f || echo 0)` is WRONG --
    # grep -c prints 0 AND exits 1, so the substitution yields $'0\n0' and
    # [ -eq ] throws. stat above is fine; it prints nothing on failure.
    ROWS=$(grep -c . "$G" 2>/dev/null); ROWS=${ROWS:-0}
    if [ "$ROWS" -ne "$N" ]; then
      say "  INCOMPLETE $NAME/$TASK ($ROWS/$N rows) -- not scoring, not marking done"
      continue
    fi
    say "  scoring $NAME / $TASK"
    $PY Tests/bluscriptcp/longeval_score.py --gen "$G" \
        --meta "$OUT/prompts/${TASK}_meta.json" \
        --out "$OUT/scores/${NAME}_${TASK}.csv" --show 2 | tee -a "$M"
    touch "$MK"
  done
done

# ---- 3. summary -------------------------------------------------------------
say "building summary"
# TASKS must come through argv: the heredoc is QUOTED (<<'PY'), so $TASKS would
# never expand inside it.
$PY - "$OUT" "$TASKS" <<'PY' | tee -a "$M"
import csv, math, os, sys, collections
out, tasks = sys.argv[1], sys.argv[2].split()
d = os.path.join(out, "scores")

# Filenames are <model>_<task>.csv and BOTH halves contain underscores, so
# fn.split("_", 1) mis-parses every arm ("arm_b_niah_single" -> "arm").
# Match the known task list from the tail instead; for/else skips non-matches so
# `name` cannot silently carry over from the previous iteration.
per_task = collections.defaultdict(dict)          # task -> model -> (hits, n)
per_depth = collections.defaultdict(dict)         # depth -> model -> (hits, n)
per_task_depth = collections.defaultdict(dict)    # (task,depth) -> model -> (hits,n)
for fn in sorted(os.listdir(d)) if os.path.isdir(d) else []:
    if not fn.endswith(".csv"):
        continue
    for task in tasks:
        if fn.endswith("_" + task + ".csv"):
            name = fn[:-(len(task) + 5)]
            break
    else:
        continue
    rows = list(csv.DictReader(open(os.path.join(d, fn))))
    if not rows:
        continue
    per_task[task][name] = (sum(int(r["ok"]) for r in rows), len(rows))
    for r in rows:
        dep = float(r["depth"])
        h, n = per_task_depth[(task, dep)].get(name, (0, 0))
        per_task_depth[(task, dep)][name] = (h + int(r["ok"]), n + 1)

if not per_task:
    print("  no scores yet"); raise SystemExit

models = sorted({m for t in per_task.values() for m in t})
pct = lambda hn: 100.0 * hn[0] / hn[1] if hn[1] else float("nan")

print("\n=== per-task accuracy %  (secondary/descriptive: ~10 per depth cell) ===")
print("  " + "task".ljust(18) + "".join(m.rjust(9) for m in models))
for task in sorted(per_task):
    print("  " + task.ljust(18) + "".join(
        (f"{pct(per_task[task][m]):8.1f}" if m in per_task[task] else "       -").rjust(9)
        for m in models))

# Degenerate-task exclusion, pre-registered: a task at <=1% overall for EVERY
# model adds n without information and would shrink the pooled CI spuriously.
degenerate = {t for t in per_task
              if all(pct(per_task[t][m]) <= 1.0 for m in per_task[t])}
if degenerate:
    print(f"\n  degenerate (<=1% for every model): {sorted(degenerate)}")
# If EVERY task is degenerate the exclusion would empty the primary table and
# report nothing at all -- a floor effect is a result, not an absence of one.
# Fall back to pooling everything and say so, rather than printing a blank table.
if len(degenerate) == len(per_task):
    print("  ALL tasks degenerate -> FLOOR EFFECT. Pooling everything anyway;")
    print("  differences between models are untestable at this setting (no power).")
    degenerate = set()

for (task, dep), byname in per_task_depth.items():
    if task in degenerate:
        continue
    for name, (h, n) in byname.items():
        ph, pn = per_depth[dep].get(name, (0, 0))
        per_depth[dep][name] = (ph + h, pn + n)

def wilson(h, n):
    """95% Wilson half-width in points -- honest near 0, unlike the normal approx."""
    if not n:
        return float("nan")
    z, p = 1.96, h / n
    return 100.0 * z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / (1 + z * z / n)

print("\n=== PRIMARY: accuracy % pooled by depth across tasks (+-95% Wilson) ===")
print("  " + "depth".ljust(8) + "".join(m.rjust(16) for m in models))
for dep in sorted(per_depth):
    cells = []
    for m in models:
        if m in per_depth[dep]:
            h, n = per_depth[dep][m]
            cells.append(f"{pct((h,n)):6.1f}+-{wilson(h,n):4.1f}".rjust(16))
        else:
            cells.append("-".rjust(16))
    print("  " + f"{dep:<8.2f}" + "".join(cells))

with open(os.path.join(out, "summary.csv"), "w", newline="") as f:
    w = csv.writer(f); w.writerow(["task"] + models)
    for task in sorted(per_task):
        w.writerow([task] + [f"{pct(per_task[task][m]):.1f}" if m in per_task[task] else ""
                             for m in models])
with open(os.path.join(out, "depth_summary.csv"), "w", newline="") as f:
    w = csv.writer(f); w.writerow(["depth", "model", "hits", "n", "pct", "wilson95"])
    for dep in sorted(per_depth):
        for m in models:
            if m in per_depth[dep]:
                h, n = per_depth[dep][m]
                w.writerow([dep, m, h, n, f"{pct((h,n)):.1f}", f"{wilson(h,n):.1f}"])
print(f"\n  wrote {out}/summary.csv and {out}/depth_summary.csv")
PY

say "DONE -> $OUT/summary.csv"