#!/bin/bash
# Launch the Phase 4 decisive seed matrix on the v2 tiny backbone: PRISME / global recurrence / dense, seeds 1..5.
# Usage: scripts/launch_seed_matrix.sh  (edits configs in place from the v2 checkpoint, syncs, launches one run per node slot)
set -e
cd "$(dirname "$0")/.."
CK=$(ls -d experiments/runs/p4_tiny_backbone_v2_seed01_*/checkpoint | tail -1)
echo "backbone checkpoint: $CK"
.venv/bin/python - "$CK" <<'PY'
import sys, yaml
ck = sys.argv[1]
model = {"name": "tiny_p4v2", "vocab_size": 128, "d_model": 256, "n_layers": 6, "n_heads": 8, "n_kv_heads": 4, "ffn_width": 768,
         "max_seq_len": 320, "w": 64, "ffn_w": 192, "n_slots": 16, "n_editor_heads": 2, "r_max": 8,
         "readout_zero_init": False, "readout_init_std": 0.2}
base = {"seed": 1, "model": model, "init": {"backbone_checkpoint": ck, "freeze_backbone": True, "auto_readout_scale": True},
        "data": {"kind": "synthetic", "tokenizer": "char", "splits_dir": "data/splits/synthetic_v1", "loss_span": "answer",
                 "eval_splits": ["dev_iid", "dev_template", "dev_depth", "dev_compose"]},
        "budgets": [0, 1, 2, 4, 8], "budget_probs": [0.2] * 5, "budget_per": "position",
        "train": {"batch_size": 32, "seq_len": 320, "steps": 3000, "lr": 1.0e-3, "warmup": 100, "weight_decay": 0.1, "grad_clip": 1.0,
                  "log_every": 50, "eval_every": 500, "eval_batches": 12, "ckpt_every": 3000},
        "em_budgets": [0, 1, 2, 8], "em_max_items": 1000}
for name, mode, extra in [("p4v2_prisme", "prisme", {}), ("p4v2_gr", "global_recurrence", {"baseline_kwargs": {"ffn_w": 640}}), ("p4v2_dense", "dense", {})]:
    e = dict(base); e["name"] = name; e["mode"] = mode; e.update(extra)
    yaml.safe_dump(e, open(f"configs/experiments/{name}.yaml", "w"), sort_keys=False)
print("configs written")
PY
# node slots: (node, config, seed)
SLOTS="M2U64:p4v2_prisme:1 M4M64c:p4v2_prisme:2 M4M64a:p4v2_prisme:3 M5M36:p4v2_prisme:4 m4mn:p4v2_prisme:5 M2U64:p4v2_gr:1 M4M64c:p4v2_gr:2 M4M64a:p4v2_gr:3 M5M36:p4v2_gr:4 m4mn:p4v2_gr:5 M2U64:p4v2_dense:1 M4M64c:p4v2_dense:2 M4M64a:p4v2_dense:3 M5M36:p4v2_dense:4 m4mn:p4v2_dense:5"
for n in M2U64 M4M64c M4M64a M5M36 m4mn; do
  rsync -az --exclude .venv --exclude /data --exclude /experiments --exclude .git ./ $n:~/wp31/prisme-slm/ &
done; wait
for n in M2U64 M4M64c M4M64a M5M36 m4mn; do rsync -az --relative ./$CK $n:~/wp31/prisme-slm/ & done; wait
for slot in $SLOTS; do
  IFS=: read -r node cfg seed <<< "$slot"
  ssh $node "cd ~/wp31/prisme-slm && nohup .venv/bin/python -m training.train_editor --config configs/experiments/$cfg.yaml --seed $seed --device mps > experiments/logs/${cfg}_seed$seed.log 2>&1 < /dev/null & disown" && echo "launched $cfg seed $seed on $node"
done