#!/usr/bin/env bash
# Batch 4 — heterogeneous colonies (Robin's mixed-size idea, PLAN v1a+).
#
# The comparison is designed so "it just has more parameters" cannot explain a win:
#   hetero-mixed   [256,64,64,64]  1.98M params  <- one elder, three juveniles
#   hetero-untied  [64,64,64,64]   0.90M params  <- same members, no elder (control)
#   solo D=256, half evidence      1.49M params  <- the elder ALONE, same blindness
# If mixed beats both, the elder is contributing something as a colony member that
# it cannot contribute alone, and the juveniles are contributing something the
# elder lacks. If mixed only beats the untied control, capacity explains it.
#
# All runs use distributed evidence (half the grid each) — the only regime where
# colonies have been shown to create capability rather than duplicate it.
set -uo pipefail
cd "$(dirname "$0")/.."

run_hetero () {  # name, widths, extra
  local name="$1" widths="$2" extra="${3:-}"
  echo "=== $(date -Is) starting ${name} ==="
  # shellcheck disable=SC2086
  PYTHONUNBUFFERED=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
  .venv/bin/python -m dsem.pretrain \
    arch=hetero \
    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.L_layers=2 arch.H_cycles=3 arch.L_cycles=6 \
    "arch.widths=${widths}" \
    arch.view_fraction=0.5 \
    ema=True "+run_name=${name}" \
    $extra > "runs/${name}.out" 2>&1 < /dev/null
  echo "=== $(date -Is) finished ${name} ==="
}

run_hetero hetero-mixed-1elder  "[256,64,64,64]"
run_hetero hetero-untied-k4     "[64,64,64,64]"
run_hetero hetero-solo-d256     "[256]"   "arch.view_enforce_coverage=False"

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