#!/usr/bin/env python3
"""Overlay intra-node and inter-node MB2 curves on the same axes."""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


def load(path: Path) -> list[dict]:
    d = json.loads(path.read_text())
    return [s for s in d["summary"] if s["input_tokens"] >= 64]


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--intra", type=Path, required=True)
    p.add_argument("--inter", type=Path, required=True)
    p.add_argument("--out-time", type=Path, default=Path("figs/mb2_transfer_time_compare.png"))
    p.add_argument("--out-bw",   type=Path, default=Path("figs/mb2_transfer_bw_compare.png"))
    args = p.parse_args()

    intra = load(args.intra)
    inter = load(args.inter)

    def axis_arrays(rows):
        kv = [r["kv_mib"] for r in rows]
        p50 = [r["pure_transfer_ms_p50"] for r in rows]
        mn = [r["pure_transfer_ms_min"] for r in rows]
        mx = [r["pure_transfer_ms_max"] for r in rows]
        bw_p50 = [r["throughput_gbps_p50"] for r in rows]
        bw_max = [r["throughput_gbps_max"] for r in rows]
        return kv, p50, mn, mx, bw_p50, bw_max

    ai_kv, ai_p50, ai_mn, ai_mx, ai_bw_p50, ai_bw_max = axis_arrays(intra)
    bi_kv, bi_p50, bi_mn, bi_mx, bi_bw_p50, bi_bw_max = axis_arrays(inter)

    # ---- transfer time ----
    fig, ax = plt.subplots(figsize=(8.5, 5))
    ax.errorbar(ai_kv, ai_p50,
                yerr=[np.array(ai_p50) - np.array(ai_mn),
                      np.array(ai_mx) - np.array(ai_p50)],
                fmt="o-", color="#1f77b4", lw=2, markersize=7, capsize=4,
                label="intra-node (dash1 GPU 0↔1)")
    ax.errorbar(bi_kv, bi_p50,
                yerr=[np.array(bi_p50) - np.array(bi_mn),
                      np.array(bi_mx) - np.array(bi_p50)],
                fmt="s--", color="#d62728", lw=2, markersize=7, capsize=4,
                label="inter-node (dash1 GPU0 → dash2 GPU0)")
    # ideal 9.7 GB/s reference
    ref_x = np.array(ai_kv)
    ref_y_ms = (ref_x * 1024 * 1024) / (9.7 * 1e9) * 1000
    ax.plot(ref_x, ref_y_ms, "--", color="#888", alpha=0.5,
            label="9.7 GB/s reference")
    ax.axvline(11500, color="#7a1d1d", lw=0.8, ls=":", alpha=0.5)
    ax.text(11500, 0.7, "p99 agentic req\n11.5 GiB",
            fontsize=8, color="#7a1d1d", ha="center")
    ax.set_xscale("log"); ax.set_yscale("log")
    ax.set_xlabel("KV transfer size (MiB)")
    ax.set_ylabel("Pure transfer time (ms, log)")
    ax.set_title("MB2 intra vs inter — Mooncake transfer cost is topology-independent\n"
                 "(both paths go through 200 Gbps RDMA NIC; intra-node does not use NVLink)")
    ax.grid(True, which="both", alpha=0.3)
    ax.legend(loc="upper left", fontsize=9)
    args.out_time.parent.mkdir(parents=True, exist_ok=True)
    fig.tight_layout(); fig.savefig(args.out_time, dpi=150); plt.close(fig)
    print(f"wrote {args.out_time}")

    # ---- bandwidth ----
    fig, ax = plt.subplots(figsize=(8.5, 5))
    ax.plot(ai_kv, ai_bw_p50, "o-", color="#1f77b4", lw=2, markersize=7,
            label="intra p50")
    ax.plot(ai_kv, ai_bw_max, "x--", color="#1f77b4", lw=1.2, markersize=8,
            alpha=0.7, label="intra max")
    ax.plot(bi_kv, bi_bw_p50, "s-", color="#d62728", lw=2, markersize=7,
            label="inter p50")
    ax.plot(bi_kv, bi_bw_max, "+--", color="#d62728", lw=1.2, markersize=8,
            alpha=0.7, label="inter max")
    ax.axhline(9.7, color="#888", ls="--", alpha=0.5,
               label="steady-state ≈ 9.7 GB/s")
    ax.set_xscale("log")
    ax.set_xlabel("KV transfer size (MiB)")
    ax.set_ylabel("Effective bandwidth (GB/s)")
    ax.set_ylim(0, 12)
    ax.set_title("MB2 intra vs inter — bandwidth")
    ax.grid(True, which="both", alpha=0.3)
    ax.legend(loc="lower left", fontsize=9)
    args.out_bw.parent.mkdir(parents=True, exist_ok=True)
    fig.tight_layout(); fig.savefig(args.out_bw, dpi=150); plt.close(fig)
    print(f"wrote {args.out_bw}")


if __name__ == "__main__":
    main()