#!/bin/bash
# Stage C (labels) -> Stage D (controller, h+entropy, early stopping) -> routing evaluation on dev-iid, for one
# finished editor run on one node. Usage: scripts/run_stage_cd.sh <node> <editor_run_id> <tag>
set -e
NODE=$1; RUN=$2; TAG=$3
cd "$(dirname "$0")/.."
.venv/bin/python - "$RUN" "$TAG" <<'PY'
import sys, yaml
run, tag = sys.argv[1], sys.argv[2]
ck = f"experiments/runs/{run}/checkpoint"
lab = yaml.safe_load(open("configs/experiments/labels_p4.yaml")); lab.update(name=f"labels_{tag}", model="from_checkpoint", checkpoint=ck)
lab["partitions"] = {"train": {"split": "ctrl_train", "n_batches": 300, "stream_seed": 11}, "calib": {"split": "calib", "n_batches": 40, "stream_seed": 12}}
yaml.safe_dump(lab, open(f"configs/experiments/labels_{tag}.yaml", "w"), sort_keys=False)
ctl = yaml.safe_load(open("configs/experiments/controller_p4_hent.yaml")); ctl.update(name=f"controller_{tag}", model="from_checkpoint", checkpoint=ck)
ctl["labels"] = {"train": f"LABELS_DIR/labels_train.pt", "calib": f"LABELS_DIR/labels_calib.pt"}
ctl["train"]["steps"] = 2000; ctl["train"]["eval_every"] = 250
yaml.safe_dump(ctl, open(f"configs/experiments/controller_{tag}.yaml", "w"), sort_keys=False)
rt = yaml.safe_load(open("configs/experiments/routing_p4.yaml")); rt["model"] = "from_checkpoint"; rt["checkpoint"] = "CONTROLLER_CKPT"
yaml.safe_dump(rt, open(f"configs/experiments/routing_{tag}.yaml", "w"), sort_keys=False)
PY
rsync -az configs training evaluation prisme baselines data/synthetic_dataset.py $NODE:~/wp31/prisme-slm/ 2>/dev/null || rsync -az configs training evaluation prisme baselines $NODE:~/wp31/prisme-slm/
ssh $NODE "cd ~/wp31/prisme-slm && nohup bash -c '
set -e
.venv/bin/python -m training.collect_interventions --config configs/experiments/labels_$TAG.yaml --device mps > experiments/logs/labels_$TAG.log 2>&1
L=\$(ls -td experiments/runs/labels_${TAG}_seed* | head -1)
sed -i \"\" \"s|LABELS_DIR|\$L|g\" configs/experiments/controller_$TAG.yaml
.venv/bin/python -m training.train_controller --config configs/experiments/controller_$TAG.yaml --device mps > experiments/logs/controller_$TAG.log 2>&1
C=\$(ls -td experiments/runs/controller_${TAG}_seed* | head -1)
sed -i \"\" \"s|CONTROLLER_CKPT|\$C/checkpoint|g\" configs/experiments/routing_$TAG.yaml
.venv/bin/python -m evaluation.routing --config configs/experiments/routing_$TAG.yaml --split dev_iid --device mps --out results/raw/routing_$TAG.json > experiments/logs/routing_$TAG.log 2>&1
echo STAGE_CD_DONE >> experiments/logs/routing_$TAG.log
' > experiments/logs/stage_cd_$TAG.log 2>&1 < /dev/null & disown; echo launched stage C/D for $RUN on $NODE"