#!/usr/bin/env python3
"""Coordinator-side distributed benchmark (Gate E): compares local bounded residency vs remote execution of a placed subset.
--workers name=host:port[,name=host:port]  --placement {none,odd,half_layers,popular:<traces>} --capacity C
Writes JSONL records (same schema as bench_decode + remote stats) and a correctness check vs mlx-lm on the first prompt.
"""
import argparse, json, os, random, socket, subprocess, sys, time, statistics
import mlx.core as mx, numpy as np, yaml
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from mlx_lm import stream_generate
from mlx_lm.utils import load_model, load_tokenizer
from titan.runtime import TitanModel
from transport.client import WorkerClient
sys.path.insert(0, str(Path(__file__).resolve().parent))
from bench_decode import run_one, env_record


def make_placement(kind, spec, workers, traces=None):
    L, E = spec.num_layers, spec.num_experts
    names = list(workers)
    if kind == "none":
        return {}
    if kind == "odd":
        return {(l, e): names[0] for l in range(L) for e in range(E) if e % 2 == 1}
    if kind == "all_remote":
        return {(l, e): names[(l * E + e) % len(names)] for l in range(L) for e in range(E)}
    if kind.startswith("popular:"):  # least popular experts go remote (coordinator keeps the hot set), split round-robin over workers
        d = kind.split(":", 1)[1]; idx = json.load(open(os.path.join(d, "index.json"))); cnt = np.zeros((L, E))
        for p in idx["prompts"]:
            if p["split"] != "fit": continue
            z = np.load(os.path.join(d, p["name"] + ".npz")); t = z["topk"]
            for l in range(L):
                cnt[l] += np.bincount(t[:, l].reshape(-1), minlength=E)
        pl = {}
        for l in range(L):
            order = np.argsort(cnt[l])  # ascending popularity
            for i, e in enumerate(order[: E // 2]):
                pl[(l, int(e))] = names[i % len(names)]
        return pl
    raise ValueError(kind)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--workers", default="")
    ap.add_argument("--token-file", default="~/.titan_token")
    ap.add_argument("--placement", default="none")
    ap.add_argument("--capacity", type=int, required=True)
    ap.add_argument("--predictor", default="none"); ap.add_argument("--m", type=int, default=8); ap.add_argument("--traces", default=None)
    ap.add_argument("--workloads", default=str(Path(__file__).resolve().parents[1] / "configs" / "workloads_v2.yaml"))
    ap.add_argument("--families", default="general,code"); ap.add_argument("--split", default="heldout")
    ap.add_argument("--max-tokens", type=int, default=64); ap.add_argument("--repeats", type=int, default=2); ap.add_argument("--warmup", type=int, default=1)
    ap.add_argument("--seed", type=int, default=0); ap.add_argument("--tag", default=""); ap.add_argument("--out", required=True)
    a = ap.parse_args()
    tm = TitanModel(a.model, capacity=a.capacity, predictor=a.predictor, m=a.m, traces_dir=a.traces)
    token = open(os.path.expanduser(a.token_file), "rb").read().strip()
    workers = {}
    for spec in [w for w in a.workers.split(",") if w]:
        name, hp = spec.split("="); host, port = hp.rsplit(":", 1)
        workers[name] = WorkerClient(host, int(port), token, tm.spec.model_type, name=name)
    placement = make_placement(a.placement, tm.spec, workers, a.traces)
    rtts = {n: w.ping(10) for n, w in workers.items()}
    if placement:
        assign = {}
        for (l, e), n in placement.items():
            assign.setdefault(n, {}).setdefault(l, []).append(e)
        for n, asg in assign.items():
            t = time.perf_counter(); r = workers[n].load(asg); print(f"worker {n}: loaded {r['loaded']} experts in {time.perf_counter()-t:.1f}s")
        rx = tm.attach_workers(workers, placement)
    else:
        rx = None
    # correctness vs mlx-lm on one prompt
    refm, _ = load_model(Path(a.model)); tok = tm.tokenizer
    ids = mx.array(tok.encode("The capital of Canada is"))[None]
    ref = refm(ids); mx.eval(ref); out = tm(ids); mx.eval(out)
    dlogit = float(mx.max(mx.abs(out.astype(mx.float32) - ref.astype(mx.float32))))
    del refm
    wl = yaml.safe_load(open(a.workloads)); prompts = [(f, p) for f in a.families.split(",") for p in wl["families"][f][a.split]]
    rng = random.Random(a.seed); order = []
    for rep in range(a.warmup + a.repeats):
        ps = prompts[:]; rng.shuffle(ps); order += [(rep, f, p) for f, p in ps]
    base = {"backend": "titan-distributed", "capacity": a.capacity, "placement": a.placement, "workers": list(workers), "rtt_ms": rtts,
            "n_remote_experts": len(placement), "predictor": a.predictor, "m": a.m, "max_tokens": a.max_tokens, "split": a.split, "tag": a.tag,
            "model": os.path.realpath(a.model), "cmd": " ".join(sys.argv), "env": env_record(), "seed": a.seed, "dlogit_vs_mlx_lm": dlogit}
    os.makedirs(os.path.dirname(a.out), exist_ok=True)
    with open(a.out, "a") as f:
        for rep, fam, p in order:
            tm.reset_stats(); mx.reset_peak_memory()
            if rx: rx.remote_pairs = rx.local_pairs = rx.layers_remote = rx.failures = 0; rx.remote_wait = 0.0
            r = run_one(tm.model, tok, p, a.max_tokens)
            cs = tm.cache_stats()
            rec = dict(base, warmup=rep < a.warmup, repeat=rep - a.warmup, family=fam, prompt=p[:60], ts=time.strftime("%Y-%m-%dT%H:%M:%S"), **r,
                       cache_hit_rate=cs["hit_rate"], cache_misses=cs["misses"], stall_s=cs["stall_time_s"], remote=rx.stats() if rx else None)
            f.write(json.dumps(rec) + "\n"); f.flush()
            print(f"dist cap={a.capacity} {a.placement:10s} rep={rep} {fam:10s} ttft={r['ttft_ms']:.0f}ms decode={r['gen_tps']:.1f} tok/s hit={cs['hit_rate']:.3f} stall={cs['stall_time_s']:.2f}s"
                  + (f" remote_wait={rx.remote_wait:.2f}s remote_pairs={rx.remote_pairs}" if rx else ""))
    for w in workers.values(): w.close()


if __name__ == "__main__":
    main()