#!/usr/bin/env python3
"""Real distributed benchmarks on MacLustr (Phase 8). The coordinator runs on `--coordinator HOST` (LAN IP), workers are launched
by ssh from this laptop on the listed nodes (several worker processes per node allowed: K partitions need not equal K hosts).

Usage:
  python3 experiments/run_benchmarks.py --config experiments/configs/benchmarks.yaml --out results/raw/benchmarks.jsonl
Each design row: model, p, N, K, splits, q, hessian_mode, repeats, workers_per_node, nodes (list of alias:lan_ip).
Requires ~/wp28/repo synced and ~/wp28/venv on every node (experiments/scatter_cluster.py does the sync).
"""
import argparse, json, pathlib, subprocess, sys, time, socket, yaml, shlex, hashlib
ROOT = pathlib.Path(__file__).resolve().parents[1]
PY = "~/wp28/venv/bin/python"


def launch_coordinator(host_alias, port, K, model, p, N, splits, q, hmode, seed, bits):
    code = (f"import sys,json; sys.path.insert(0,'{'~/wp28/repo/src'}'); import os; sys.path[0]=os.path.expanduser(sys.path[0]); "
            f"from pamle.protocol import coordinator_run; "
            f"r=coordinator_run({port},{K},'{model}',{p},{N},{tuple(splits)},{q},'{hmode}',{seed},{bits},log=lambda *a: None); print(json.dumps(r, default=lambda o: o.item() if hasattr(o, 'item') else str(o)))")
    cmd = f"{PY} -c {shlex.quote(code)}"
    return subprocess.Popen(["ssh", "-o", "ConnectTimeout=10", host_alias, cmd], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)


def launch_workers(assign, coord_ip, port):
    procs = []
    wid = 0
    for alias, n in assign:
        # workers are children of the ssh shell and `wait` keeps the session (and the workers) alive until they exit
        cmds = "cd ~/wp28/repo/src && " + " ".join(f"{PY} -m pamle.worker --connect {coord_ip}:{port} --id {wid + i} > /tmp/wp28_worker_{wid + i}.log 2>&1 &" for i in range(n)) + " wait"
        wid += n
        procs.append(subprocess.Popen(["ssh", "-o", "ConnectTimeout=10", alias, cmds], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True))
    return procs


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--config", required=True); ap.add_argument("--out", required=True); ap.add_argument("--only", default=None)
    a = ap.parse_args(); cfg = yaml.safe_load(open(a.config)); out = pathlib.Path(a.out); out.parent.mkdir(parents=True, exist_ok=True)
    coord_alias, coord_ip = cfg["coordinator"]["alias"], cfg["coordinator"]["lan_ip"]
    port = int(cfg.get("port", 50600))
    for i, d in enumerate(cfg["designs"]):
        if a.only and str(i) not in a.only.split(","):
            continue
        nodes = d["nodes"]; wpn = d["workers_per_node"]; K = d["K"]
        assign = []; left = K
        for alias in nodes:
            n = min(wpn, left); assign.append((alias, n)); left -= n
            if left == 0:
                break
        if left:
            print(f"design {i}: not enough worker slots for K={K}", file=sys.stderr); continue
        for rep in range(d.get("repeats", 3)):
            port += 1
            seed = int(hashlib.sha256(f"{i}-{rep}".encode()).hexdigest(), 16) % 10 ** 6
            co = launch_coordinator(coord_alias, port, K, d["model"], d["p"], d["N"], d["splits"], d["q"], d["hessian_mode"], seed, d.get("bits"))
            time.sleep(1.5)
            t0 = time.time(); wk = launch_workers(assign, coord_ip, port)
            try:
                sout, serr = co.communicate(timeout=1800)
            except subprocess.TimeoutExpired:
                co.kill(); print(f"design {i} rep {rep}: TIMEOUT", file=sys.stderr); continue
            for w in wk:
                w.wait(timeout=60)
            if co.returncode != 0 or not sout.strip():
                print(f"design {i} rep {rep}: coordinator failed: {serr[-800:]}", file=sys.stderr); continue
            res = json.loads(sout.strip().splitlines()[-1])
            res.update({"design_index": i, "repeat": rep, "launch_wall_s": time.time() - t0, "assignment": assign, "label": d.get("label", "")})
            with open(out, "a") as f:
                f.write(json.dumps(res) + "\n")
            st = res["stats"]
            print(f"design {i} rep {rep}: K={K} hosts={res['n_hosts']} p={d['p']} q={d['q']} {d['hessian_mode']} | fit+upload {st['stages']['A_B_local_fit_and_upload']['wall']:.3f}s "
                  f"| rounds {[round(v['wall'],3) for k,v in st['stages'].items() if k.startswith('D_F')]} | up {st['bytes_up']} B down {st['bytes_down']} B | total {res['t_total_after_connect']:.3f}s", flush=True)


if __name__ == "__main__":
    main()