#!/usr/bin/env bash
# Launch a detached colony run.
#
#   scripts/launch_colony.sh <run_name> <width> <K> [extra hydra overrides...]
#
# Same protocol as the solo baselines (50k epochs, batch 768, subsampled evals),
# so colony results are directly comparable to the capacity-floor curve:
#   D=512 84.8% | D=256 67.8% | D=128 38.8% | D=64 13.6%
set -euo pipefail

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

RUN_NAME="${1:?usage: launch_colony.sh <run_name> <width> <K> [overrides...]}"
WIDTH="${2:?need width}"
K="${3:?need K}"
shift 3 || true

mkdir -p runs
if [ -f "runs/${RUN_NAME}.pid" ] && kill -0 "$(cat "runs/${RUN_NAME}.pid")" 2>/dev/null; then
  echo "run ${RUN_NAME} already alive" >&2; exit 1
fi

PYTHONUNBUFFERED=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
nohup .venv/bin/python -m dsem.pretrain \
  arch=colony \
  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=${WIDTH}" \
  "arch.num_agents=${K}" \
  ema=True \
  "+run_name=${RUN_NAME}" \
  "$@" > "runs/${RUN_NAME}.out" 2>&1 < /dev/null &

echo $! > "runs/${RUN_NAME}.pid"
echo "launched ${RUN_NAME} (D=${WIDTH}, K=${K}) pid $(cat "runs/${RUN_NAME}.pid")"