#!/usr/bin/env python3
"""Distribute Monte Carlo cells over MacLustr nodes (greedy by estimated cost, weighted by worker slots).

Writes experiments/cluster_assignment.json, syncs the repo to each node, and launches `xargs -P` worker pools.
Usage: python3 experiments/scatter_cluster.py --configs a.yaml b.yaml --launch
Collect later with experiments/collect_cluster.sh.
"""
import argparse, json, os, pathlib, subprocess, sys
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1] / "experiments"))
import yaml
from run_monte_carlo import expand_cells

ROOT = pathlib.Path(__file__).resolve().parents[1]
NODES = {"M3U96a": 6, "M3U96b": 6, "M2U64": 4, "M4M64a": 4, "M4M64b": 4, "M4M64c": 4, "M5M36": 4, "M4BP48": 3, "M4BP36": 3,
         "M2M32c": 3, "m4mm": 3, "m4mn": 3, "m6ma": 3, "m6mb": 3, "m5ma": 3, "m4md": 2, "m4ml": 2, "M3BA24": 2}


def cost(c):
    extras = 1.6 if c["extras"] else 1.0
    return c["N"] * (c["p"] ** 1.5) * c["reps"] * (1 + 0.5 * c["splits_max"]) * extras


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--configs", nargs="+", required=True); ap.add_argument("--launch", action="store_true")
    a = ap.parse_args()
    jobs = []
    for cfg in a.configs:
        cells = expand_cells(yaml.safe_load(open(cfg)))
        for i, c in enumerate(cells):
            jobs.append({"config": os.path.relpath(cfg, ROOT), "cell": i, "cost": cost(c), "name": f'{c["label"]}_p{c["p"]}_N{c["N"]}_K{c["K"]}'})
    jobs.sort(key=lambda j: -j["cost"])
    load = {n: 0.0 for n in NODES}; assign = {n: [] for n in NODES}
    for j in jobs:
        n = min(NODES, key=lambda k: load[k] / NODES[k])
        assign[n].append(j); load[n] += j["cost"]
    total = sum(j["cost"] for j in jobs)
    print(f"{len(jobs)} jobs; est. share per node (%):", {n: round(100 * load[n] / total, 1) for n in NODES})
    json.dump({"nodes": NODES, "assignment": assign}, open(ROOT / "experiments" / "cluster_assignment.json", "w"), indent=1)
    if not a.launch:
        return
    for n, js in assign.items():
        if not js:
            continue
        script = "\n".join(f'{j["config"]} {j["cell"]}' for j in js) + "\n"
        (ROOT / "logs" / f"jobs_{n}.txt").write_text(script)
        subprocess.run(["rsync", "-az", "--exclude", "results", "--exclude", "logs", "--exclude", "references", "--exclude", "paper",
                        "--exclude", ".git", f"{ROOT}/", f"{n}:~/wp28/repo/"], check=True)
        subprocess.run(["scp", "-q", str(ROOT / "logs" / f"jobs_{n}.txt"), f"{n}:~/wp28/jobs.txt"], check=True)
        W = NODES[n]
        cmd = (f"cd ~/wp28/repo && mkdir -p logs results/raw && nohup sh -c 'cat ~/wp28/jobs.txt | xargs -P {W} -L 1 sh -c "
               f"\"OMP_NUM_THREADS=2 ~/wp28/venv/bin/python experiments/run_monte_carlo.py --config \\$0 --cells \\$1 "
               f">> logs/cell_\\$(basename \\$0 .yaml)_\\$1.log 2>&1\"' > logs/pool.log 2>&1 &")
        subprocess.run(["ssh", n, cmd], check=True)
        print(f"launched {len(js)} jobs on {n} with {W} workers")


if __name__ == "__main__":
    main()