#!/usr/bin/env bash
# =============================================================================
# run_peer_eval.sh -- NIAH on peer base checkpoints, scored identically to ours.
#
# WHAT THIS ANSWERS
#   Our 1.06B scores 48% on niah_single at 16k. Every published comparison is
#   either the wrong metric (RULER 13-task average, not single-NIAH), the wrong
#   checkpoint type (instruct, not base), or the wrong size. Running the peers
#   through OUR harness removes all three mismatches at once.
#
# MODELS (all BASE checkpoints)
#   Qwen3-0.6B/1.7B/4B/8B-Base   dense attention, native 32,768 ctx, no YaRN.
#                                The ladder gives a measured retrieval-vs-scale
#                                curve without training anything. Qwen3-1.7B is
#                                the closest peer to our 1.06B (1.41B non-emb
#                                vs our 0.906B); Qwen3-4B matches our 4B target.
#   Spark-X2.5-4B-Base           hybrid: 27 sliding-window(512) + 9 full-attention
#                                layers, claims NATIVE 1M context, publishes no
#                                NIAH/RULER number at all. custom_code, so vLLM
#                                cannot load it -- peer_run.py falls back to hf.
#
# FAIRNESS NOTES, all of which belong in any writeup
#   * One shared haystack TEXT; each model cut to its own exact token count.
#   * Qwen3/Spark run INSIDE their native windows at 16k. Ours is YaRN-extended
#     from 4k and at 16k sits at the very edge. Not the same condition.
#   * Peers use a KV cache; our CP_GENERATE re-runs a full forward per token.
#     That is an engine difference, not a capability difference.
#
# usage:  ./run_peer_eval.sh            # all models, all lengths
#         MODELS=qwen3-4b LENS=16384 ./run_peer_eval.sh
# =============================================================================
set -uo pipefail
cd "$(dirname "$0")"

PY=${PY:-python3}
ROOT=${ROOT:-$PWD/peer_out}
HAY=${HAY:-$ROOT/haystack.txt}
LENS=${LENS:-"1024 2048 4096 8192 16384"}
N=${N:-100}
TASKS=${TASKS:-"niah_single val_backward"}
MAX_NEW=${MAX_NEW:-24}
DEPTHS=${DEPTHS:-"0.0,0.25,0.5,0.75,1.0"}
GPUS=${GPUS:-"0 1 2 3 4 5 6 7"}
MODELS=${MODELS:-"qwen3-0.6b qwen3-1.7b qwen3-4b qwen3-8b spark-4b"}

# Spark-X2.5 needs transformers 4.57.x: modeling_spark.py targets 4.57 and breaks
# on 5.x twice over -- _tied_weights_keys became a mapping, and create_causal_mask
# lost its `input_embeds` kwarg. Rather than shim each API change, give it the
# version it was written for. Qwen3 stays on the vLLM venv.
PY_SPARK=${PY_SPARK:-/home/blubridge/venvs/tf457/bin/python}
declare -A PYFOR=( [spark-4b]="$PY_SPARK" )

declare -A REPO=(
  [qwen3-0.6b]=Qwen/Qwen3-0.6B-Base
  [qwen3-1.7b]=Qwen/Qwen3-1.7B-Base
  [qwen3-4b]=Qwen/Qwen3-4B-Base
  [qwen3-8b]=Qwen/Qwen3-8B-Base
  [spark-4b]=XHToken/Spark-X2.5-4B-Base
)

mkdir -p "$ROOT"
M="$ROOT/master.log"
say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$M" >&2; }

say "=== peer NIAH eval ==="
for m in $MODELS; do [ -n "${REPO[$m]:-}" ] || { say "unknown model key: $m"; exit 1; }; done

# ---- 1. one shared haystack ------------------------------------------------
if [ ! -s "$HAY" ]; then
  say "building shared haystack"
  $PY make_haystack.py --out "$HAY" --tokens 4000000 >>"$M" 2>&1 \
    || { say "FATAL: haystack build failed"; exit 1; }
fi
say "haystack $(stat -Lc%s "$HAY") bytes"

# ---- 2. prompts, per model (tokenizer differs) and per length --------------
for m in $MODELS; do
  for L in $LENS; do
    d="$ROOT/prompts/$m/$L"
    # Must check EVERY task, not just the first: an interrupt between tasks would
    # otherwise leave this dir looking complete and the generation step would fail
    # on the missing prompts file.
    have=1; for t in $TASKS; do [ -f "$d/${t}_meta.json" ] || have=0; done
    [ "$have" = 1 ] && { say "prompts exist $m/$L"; continue; }
    mkdir -p "$d"
    for t in $TASKS; do
      $PY peer_gen.py --tokenizer "${REPO[$m]}" --haystack "$HAY" --task "$t" \
          --prompt-len "$L" --n "$N" --depths "$DEPTHS" --seed 0 --out-dir "$d" \
          >>"$M" 2>&1 || { say "FATAL: prompt gen failed $m/$L/$t"; exit 1; }
    done
    say "prompts $m/$L"
  done
done

# ---- 3. generate: a real work queue, not round-robin ----------------------
# The previous version picked the GPU by LAUNCH INDEX (i % NG). `wait -n` kept
# NG jobs in flight, but the card was chosen by counter rather than by which one
# was actually free -- and a 16,384-token job takes ~16x a 1,024-token one, so
# the counter desynchronised from real occupancy within seconds and left half
# the box idle. Now: one persistent worker per GPU, each pulling the next job
# from a shared cursor under flock, so a card is never idle while work remains.
#
# Jobs are ordered LONGEST FIRST (16384 -> 1024). Longest-processing-time-first
# is the standard greedy makespan heuristic; scheduling the big jobs last would
# end the run with one 16k job while seven GPUs sit idle.
JOBS="$ROOT/joblist.txt"; : > "$JOBS"
for L in $(echo $LENS | tr ' ' '\n' | sort -rn); do
  for m in $MODELS; do
    for t in $TASKS; do
      out="$ROOT/gen/$L/${m}_${t}.jsonl"
      # require the FULL row count: a job killed mid-write would otherwise look
      # done to a `-s` test and silently truncate that cell.
      if [ -f "$out" ] && [ "$(wc -l < "$out")" -eq "$N" ]; then continue; fi
      mkdir -p "$(dirname "$out")"
      echo "$m|$L|$t" >> "$JOBS"
    done
  done
done
NJOB=$(grep -c . "$JOBS" 2>/dev/null || echo 0)
say "queue: $NJOB job(s) to run on $(echo $GPUS | wc -w) GPU(s), longest first"

if [ "$NJOB" -gt 0 ]; then
  CUR="$ROOT/.cursor"; echo 0 > "$CUR"
  worker(){
    local g="$1" idx line m L t
    while :; do
      idx=$(flock "$CUR" -c 'n=$(cat '"$CUR"'); echo $((n+1)) > '"$CUR"'; echo $n')
      line=$(sed -n "$((idx+1))p" "$JOBS")
      [ -n "$line" ] || break
      m=${line%%|*}; rest=${line#*|}; L=${rest%%|*}; t=${rest##*|}
      say "  [gpu $g] $m / $L / $t"
      CUDA_VISIBLE_DEVICES=$g "${PYFOR[$m]:-$PY}" peer_run.py --model "${REPO[$m]}" \
        --prompts "$ROOT/prompts/$m/$L/${t}_prompts.jsonl" \
        --out "$ROOT/gen/$L/${m}_${t}.jsonl" --max-new "$MAX_NEW" \
        > "$ROOT/gen/$L/${m}_${t}.log" 2>&1 \
        || say "  FAIL [gpu $g] $m/$L/$t (see $ROOT/gen/$L/${m}_${t}.log)"
    done
  }
  for g in $GPUS; do worker "$g" & done
  wait
fi
say "generation done"

# ---- 4. score, one summary per length --------------------------------------
CSV=$(echo $MODELS | tr ' ' ','); TSV=$(echo $TASKS | tr ' ' ',')
for L in $LENS; do
  say "scoring $L"
  $PY peer_score.py --gen-dir "$ROOT/gen/$L" --meta-dir "$ROOT/prompts" \
      --out-dir "$ROOT/results/$L" --length "$L" --models "$CSV" --tasks "$TSV" \
      >>"$M" 2>&1 || say "  WARN: scoring failed at $L"
  [ -f "$ROOT/results/$L/summary.csv" ] && say "  -> $(sed -n 2p "$ROOT/results/$L/summary.csv")"
done
say "=== PEER EVAL DONE -> $ROOT/results ==="