"""3-way GPU utilization + latency comparison."""
import csv, json, statistics, os

def gpu_stats(path, groups):
    rows = list(csv.DictReader(open(path)))
    by_gpu = {}
    for r in rows:
        g = int(r["gpu"])
        by_gpu.setdefault(g, []).append(float(r["util_pct"]))
    result = {}
    for gname, indices in groups.items():
        vals = []
        for i in indices:
            vals.extend(by_gpu.get(i, []))
        if vals:
            s = sorted(vals)
            p = lambda q: s[min(int(q*len(s)), len(s)-1)]
            nz = sum(1 for v in vals if v > 0)
            result[gname] = {
                "mean": statistics.fmean(vals), "p50": p(.5), "p90": p(.9),
                "max": max(vals), "active_pct": nz*100//len(vals), "n": len(rows)//8
            }
    return result

def latency_stats(path):
    rows = [json.loads(l) for l in open(path)]
    ok = [r for r in rows if not r.get("error")]
    ttfts = sorted([r["ttft_s"] for r in ok if r.get("ttft_s")])
    tpots = sorted([r["tpot_s"] for r in ok if r.get("tpot_s") and r["tpot_s"]>0])
    lats = sorted([r["latency_s"] for r in ok])
    p = lambda v,q: v[min(int(q*len(v)),len(v)-1)] if v else 0
    return {
        "ok": len(ok), "total": len(rows),
        "ttft50": p(ttfts,.5), "ttft90": p(ttfts,.9),
        "tpot50": p(tpots,.5), "tpot90": p(tpots,.9),
        "e2e50": p(lats,.5), "e2e90": p(lats,.9),
    }

configs = [
    ("gpu_ab_combined", "Combined 8colo", {"All(0-7)": list(range(8))}),
    ("gpu_ab_pdsep",    "PD-Sep 4P+4D",  {"P(0-3)": [0,1,2,3], "D(4-7)": [4,5,6,7], "All": list(range(8))}),
    ("gpu_ab_6p2d",     "PD-Sep 6P+2D",  {"P(0-5)": list(range(6)), "D(6-7)": [6,7], "All": list(range(8))}),
]

sep = "=" * 85
print(sep)
print("  3-WAY COMPARISON: Combined vs 4P+4D vs 6P+2D")
print("  Same cache-aware scheduler, 200 req, time_scale=20")
print(sep)

print("\n  GPU UTILIZATION:")
fmt = "  %-22s %7s %7s %7s %7s %7s"
print(fmt % ("Config / Group", "Mean%", "P50%", "P90%", "Max%", "Active"))
print("  " + "-" * 62)

for dirname, label, groups in configs:
    gpath = "outputs/%s/gpu_util.csv" % dirname
    if not os.path.exists(gpath):
        continue
    gs = gpu_stats(gpath, groups)
    for gname, s in gs.items():
        tag = "%s %s" % (label, gname)
        print(fmt % (tag[:22], "%.1f" % s["mean"], "%.0f" % s["p50"],
                      "%.0f" % s["p90"], "%.0f" % s["max"], "%d%%" % s["active_pct"]))

print("\n  LATENCY:")
fmt2 = "  %-17s %6s %8s %8s %8s %8s %8s"
print(fmt2 % ("Config", "OK/N", "TTFT50", "TTFT90", "TPOT50", "TPOT90", "E2E50"))
print("  " + "-" * 65)

for dirname, label, _ in configs:
    mpath = "outputs/%s/metrics.jsonl" % dirname
    if not os.path.exists(mpath):
        continue
    s = latency_stats(mpath)
    okn = "%d/%d" % (s["ok"], s["total"])
    print(fmt2 % (label[:17], okn, "%.3f" % s["ttft50"], "%.3f" % s["ttft90"],
                   "%.3f" % s["tpot50"], "%.3f" % s["tpot90"], "%.3f" % s["e2e50"]))