#!/usr/bin/env bash
# Capacity-floor sweep: solo TRMs at decreasing hidden width, same protocol as the
# reproduction run, so every point is comparable to the D=512 anchor (84.5%).
#
#   scripts/run_sweep.sh [widths...]        # default: 256 128 64
#
# Runs sequentially in one detached process; each width gets its own run name,
# log, metrics and checkpoints. Progress: runs/sweep.out
#
# This curve is the x-axis of the whole DSEM project: it says how badly a lone
# unit of each size does, which is what a colony of such units must beat.
set -euo pipefail

cd "$(dirname "$0")/.."

WIDTHS=("$@")
[ ${#WIDTHS[@]} -eq 0 ] && WIDTHS=(256 128 64)

for W in "${WIDTHS[@]}"; do
  RUN="sweep-d${W}"
  echo "=== $(date -Is) starting ${RUN} ==="
  PYTHONUNBUFFERED=1 .venv/bin/python -m dsem.pretrain \
    arch=trm \
    data_paths="[data/sudoku-extreme-1k-aug-1000]" \
    data_paths_test="[data/sudoku-testsub-12k]" \
    evaluators="[]" \
    epochs=50000 eval_interval=5000 \
    global_batch_size=768 \
    lr=1e-4 puzzle_emb_lr=1e-4 weight_decay=1.0 puzzle_emb_weight_decay=1.0 \
    arch.mlp_t=True arch.pos_encodings=none \
    arch.L_layers=2 arch.H_cycles=3 arch.L_cycles=6 \
    "arch.hidden_size=${W}" \
    ema=True \
    "+run_name=${RUN}" > "runs/${RUN}.out" 2>&1 < /dev/null
  echo "=== $(date -Is) finished ${RUN} ==="
done

echo "=== $(date -Is) sweep complete ==="