#!/usr/bin/env bash
# One-click ZBP scaling ladder. Clone -> ./scripts/run_ladder.sh   (run on the H200 node, nohup-able)
#
# Configurable via env vars (defaults in brackets):
#   SIZES="s60 m124 l350"        model configs from configs/model/   [s60 m124 l350]
#   ARMS="bp zbp_n16 zbp_n64"    train configs from configs/train/   [bp zbp_n16 zbp_n64]
#   NPROC=8                      GPUs per run                        [auto: nvidia-smi -L]
#   TOKENS=2e9                   override token budget per run       [per train config]
#   DATA=data/fineweb            token shards (prepared if missing)  [data/fineweb]
#   PREP_TOKENS=3e9              tokens to tokenize if DATA missing  [3e9]
#   OUT=runs                     output root                        [runs]
#   MICRO_BS=8                   per-GPU micro batch                [8]
#   SET="k=v k2=v2"              extra --set overrides for every run (e.g. vocab=8192 seq_len=256)
#   DRY=1                        print the plan and exit
#   PUSH_RESULTS=0               skip the auto collect + git-push of results JSONL   [on]
#   HF_UPLOAD=1 [HF_CKPT=1]      also upload to a private HF repo (needs HF_TOKEN; see README)
#
# Runs are sequential (each takes the whole node), resume-safe: a finished run leaves OUT/<name>/DONE
# and is skipped on re-invocation, so the script can be re-run after interruptions.
set -euo pipefail
cd "$(dirname "$0")/.."

SIZES=${SIZES:-"s60 m124 l350"}
ARMS=${ARMS:-"bp zbp_n16 zbp_n64"}
NPROC=${NPROC:-$(nvidia-smi -L 2>/dev/null | wc -l)}
DATA=${DATA:-data/fineweb}
PREP_TOKENS=${PREP_TOKENS:-3e9}
OUT=${OUT:-runs}
MICRO_BS=${MICRO_BS:-8}
EXTRA="--set ${TOKENS:+tokens=$TOKENS} ${SET:-}"
[ "$EXTRA" = "--set  " ] && EXTRA=""

echo "== zbp-scaling ladder: sizes=[$SIZES] arms=[$ARMS] nproc=$NPROC data=$DATA out=$OUT ${TOKENS:+tokens=$TOKENS}"
[ "$NPROC" -ge 1 ] || { echo "no GPUs detected; set NPROC"; exit 1; }
# --- environment bootstrap: current env first; VENV=1 creates ./.venv (fresh nodes) ---
PY=${PY:-python3}
if [ "${VENV:-0}" = 1 ]; then
  [ -d .venv ] || "$PY" -m venv .venv
  . .venv/bin/activate
fi
python -c "import zbp_scaling" 2>/dev/null || pip install -e . -q || {
  echo "pip install failed; on managed nodes run with VENV=1 (creates ./.venv) or install into your conda env"; exit 1; }
python scripts/env_check.py || exit 1

if [ ! -f "$DATA/train.bin" ]; then
  echo "== preparing data ($PREP_TOKENS tokens -> $DATA) — run this on the training node"
  [ "${DRY:-0}" = 1 ] || python scripts/prepare_data.py --tokens "$PREP_TOKENS" --out "$DATA"
fi

for size in $SIZES; do
  for arm in $ARMS; do
    [ -f "configs/model/$size.yaml" ] || { echo "missing configs/model/$size.yaml"; exit 1; }
    [ -f "configs/train/$arm.yaml" ] || { echo "missing configs/train/$arm.yaml"; exit 1; }
    name="${size}_${arm}"
    dir="$OUT/$name"
    if [ -f "$dir/DONE" ]; then echo "== skip $name (done)"; continue; fi
    echo "== run $name  (model=configs/model/$size.yaml train=configs/train/$arm.yaml)"
    [ "${DRY:-0}" = 1 ] && continue
    mkdir -p "$dir"
    if [ "$NPROC" -gt 1 ]; then
      torchrun --standalone --nproc_per_node="$NPROC" scripts/train.py \
        --model "configs/model/$size.yaml" --train "configs/train/$arm.yaml" \
        --data "$DATA" --out "$dir" --micro_bs "$MICRO_BS" $EXTRA 2>&1 | tee -a "$dir/stdout.log"
    else
      python scripts/train.py --model "configs/model/$size.yaml" --train "configs/train/$arm.yaml" \
        --data "$DATA" --out "$dir" --micro_bs "$MICRO_BS" $EXTRA 2>&1 | tee -a "$dir/stdout.log"
    fi
    touch "$dir/DONE"
  done
done
TAG=${HF_TAG:-$(hostname)-$(date +%Y%m%d)}
if [ "${PUSH_RESULTS:-1}" = 1 ] && [ "${DRY:-0}" != 1 ]; then
  # zero-token default: collect the small JSONL results and push them back over this clone's git auth
  if python scripts/collect.py --runs "$OUT" --out "results/$TAG"; then
    git add results
    git -c user.email=ladder@zbp -c user.name=ladder commit -m "results: $TAG" || true  # nothing new is fine
    git push || echo "!! git push failed — push manually later or send results/$TAG"
  else
    echo "!! collect found no finished runs — skipping push"
  fi
fi
if [ "${HF_UPLOAD:-0}" = 1 ]; then
  # optional direct-to-HF (needs HF_TOKEN in the environment; see README Security)
  python scripts/collect.py --runs "$OUT" --out "results/$TAG"
  python scripts/upload_hf.py --results "results/$TAG" ${HF_REPO:+--repo "$HF_REPO"} ${HF_CKPT:+--with-ckpt "$OUT"}
fi
echo "== ladder complete: $OUT"