#!/usr/bin/env python3
"""Benchmark figures: (a) exchange rate — local phase vs synchronized rounds by protocol at p = 8 and p = 64 (16 hosts);
(b) hosts scaling of the (2,1) protocol at K = 16 workers.  Source: results/processed/benchmarks_summary.csv."""
import sys, pathlib
import numpy as np, pandas as pd
import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1] / "experiments")); from process_results import _setup, PAL, INK, INK2, FIG, PROC
_setup(); g = pd.read_csv(PROC / "benchmarks_summary.csv")
order8 = [("cmp_r4q0", "(4,0)"), ("cmp_r3q1", "(3,1)"), ("hosts16", "(2,1)"), ("cmp_r2q1_stale", "(2,1) stale"), ("cmp_r2q1_single", "(2,1) single"), ("cmp_r1q2", "(1,2)")]
order32 = [("p32_r3q1", "(3,1)"), ("p32_r2q1", "(2,1)"), ("p32_r2q1_stale", "(2,1) stale"), ("p32_r1q2", "(1,2)")]
order64 = [("cmp_p64_r2q1", "(2,1)"), ("cmp_p64_r2q1_stale", "(2,1) stale"), ("cmp_p64_r1q2", "(1,2)")]
has32 = g.label.isin([k for k, _ in order32]).any()
panels = [(order8, "p = 8, 16 hosts, K = 16")] + ([(order32, "p = 32, 16 hosts, K = 16")] if has32 else []) + [(order64, "p = 64, 16 hosts, K = 16")]
fig, axes = plt.subplots(1, len(panels) + 1, figsize=(2.4 * len(panels) + 4.4, 3.0), gridspec_kw={"width_ratios": [2.2] + [1.5] * (len(panels) - 1) + [1.7]})
for ax, (order, title) in zip(axes[:-1], panels):
    labs, loc, rnd = [], [], []
    for key, lab in order:
        r = g[g.label == key]
        if r.empty:
            continue
        labs.append(lab); loc.append(float(r.t_local.iloc[0])); rnd.append(float(r.t_rounds.iloc[0]))
    x = np.arange(len(labs)); w = 0.62
    ax.bar(x, loc, w, color=PAL["blue"], label="local phase (slowest worker)", edgecolor="white", linewidth=1.5)
    ax.bar(x, rnd, w, bottom=loc, color=PAL["orange"], label="synchronized derivative phases", edgecolor="white", linewidth=1.5)
    for xi, (a, b) in enumerate(zip(loc, rnd)):
        ax.text(xi, a + b, f"{a + b:.2f} s", ha="center", va="bottom", fontsize=7, color=INK)
    ax.set_xticks(x); ax.set_xticklabels(labs, rotation=20, ha="right", fontsize=7.5); ax.set_ylabel("wall-clock (s)"); ax.set_title(title, fontsize=8, loc="left")
    ax.set_ylim(0, 1.18 * max(a + b for a, b in zip(loc, rnd)))
axes[0].legend(fontsize=7, loc="upper center", bbox_to_anchor=(0.5, -0.32), ncol=1)
ax = axes[-1]; h = g[g.label.str.startswith("hosts")].sort_values("hosts")
ax.plot(h.hosts, h.t_local, marker="o", color=PAL["blue"], label="local phase", markerfacecolor="white", markeredgewidth=1.2)
ax.plot(h.hosts, h.fit_mean, marker="s", color=PAL["aqua"], ls="--", label="mean worker fit", markerfacecolor="white", markeredgewidth=1.2)
ax.plot(h.hosts, h.t_rounds, marker="^", color=PAL["orange"], label="derivative phase", markerfacecolor="white", markeredgewidth=1.2)
ax.set_xscale("log", base=2); ax.set_xticks([1, 2, 4, 8, 16]); ax.set_xticklabels(["1", "2", "4", "8", "16"]); ax.set_xlabel("hosts (K = 16 workers)"); ax.set_ylabel("seconds"); ax.set_title("(2,1), p = 8: host scaling", fontsize=8, loc="left"); ax.legend(fontsize=7)
FIG.mkdir(parents=True, exist_ok=True); fig.savefig(FIG / "fig_bench_exchange.pdf"); print("ok")