"""Plot paired raw samples without pooling campaigns or overlapping memory counters."""

import json
from pathlib import Path

import matplotlib

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


def draw(path, variant, title, destination):
    data = json.loads(path.read_text())
    fig, axes = plt.subplots(1, 2, figsize=(10.2, 4.5), layout="constrained")
    colors = ["#62738A", "#087F8C"]
    labels = ["Paired current control", "Exact Q/K + FC2 tuning"]
    for j, case in enumerate(["product", "portrait"]):
        for k, (offset, name) in enumerate([(-0.18, "control"), (0.18, variant)]):
            row = data["cases"][case][name]
            x = j + offset
            med = row["median_s"]
            axes[0].bar(
                x, med, width=0.32, color=colors[k], label=labels[k] if j == 0 else None
            )
            axes[0].scatter(
                [x] * len(row["samples_s"]),
                row["samples_s"],
                s=15,
                color="black",
                zorder=3,
            )
            axes[0].text(x, med + 1.7, f"{med:.2f}", ha="center", fontsize=9)
            peak = row["peak_physical_bytes"] / 1024**3
            axes[1].bar(x, peak, width=0.32, color=colors[k])
            axes[1].text(x, peak + 0.5, f"{peak:.2f}", ha="center", fontsize=9)
    for ax in axes:
        ax.set_xticks([0, 1], ["Product 2048²", "Portrait 2048²"])
        ax.spines[["top", "right"]].set_visible(False)
        ax.set_axisbelow(True)
        ax.grid(axis="y", alpha=0.2)
    axes[0].set_ylabel("Complete warm generation (s)")
    axes[0].set_ylim(
        0,
        max(
            r[name]["max_s"]
            for r in data["cases"].values()
            for name in ["control", variant]
        )
        + 8,
    )
    axes[0].set_title("Three paired repeats; dots are raw samples")
    axes[1].set_ylabel("Maximum sampled parent physical footprint (GiB)")
    axes[1].set_ylim(0, 50)
    axes[1].set_title("Includes Metal; 100 ms sampling; counters not summed")
    fig.legend(
        *axes[0].get_legend_handles_labels(),
        loc="outside lower center",
        ncol=2,
        fontsize=8,
    )
    fig.suptitle(title + " · M5 Pro / MLX 0.32.4", fontsize=12)
    fig.savefig(destination, dpi=170)
    plt.close(fig)


if __name__ == "__main__":
    root = Path("reports/optimization-r3")
    for folder, title, output in [
        ("qk-deployment", "Native BF16 weight storage", "timing-memory.png"),
        (
            "compact-deployment",
            "Lossless packed BF16 weight storage",
            "compact-timing-memory.png",
        ),
    ]:
        path = root / folder / "summary.json"
        if path.exists():
            draw(path, "qk-fc2-native-norm", title, root / output)