"""Render recorded sample timings; requires matplotlib (no model weights)."""

import json
from pathlib import Path

import matplotlib

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


def main():
    root = Path(__file__).resolve().parents[1]
    plt.rcParams.update(
        {
            "font.family": "DejaVu Sans",
            "font.size": 11,
            "svg.fonttype": "none",
            "axes.spines.top": False,
            "axes.spines.right": False,
        }
    )
    fig, axes = plt.subplots(1, 2, figsize=(11.6, 4.5), sharey=True)
    fig.patch.set_facecolor("#f7fafc")
    colors = ["#8296ac", "#137f82"]
    for ax, storage, title in zip(
        axes, ["bf16", "packed12"], ["BF16", "Lossless packed BF16"], strict=True
    ):
        ax.set_facecolor("#f7fafc")
        summary = json.loads(
            (root / f"benchmarks/data/{storage}/summary.json").read_text()
        )
        comparison = json.loads(
            (root / f"benchmarks/data/{storage}/comparison.json").read_text()
        )
        for case_index, case in enumerate(["product", "portrait"]):
            for index, (variant, label) in enumerate(
                [("control", "Control"), ("qk-fc2-native-norm", "Tuned")]
            ):
                values = summary["cases"][case][variant]["samples_s"]
                x = case_index + (-0.16 if index == 0 else 0.16)
                ax.scatter(
                    x + np.array([-0.045, 0, 0.045]),
                    values,
                    color=colors[index],
                    s=35,
                    zorder=3,
                    label=label if case_index == 0 else None,
                )
                median = float(np.median(values))
                ax.plot(
                    [x - 0.11, x + 0.11], [median, median], color=colors[index], lw=2.3
                )
                ax.annotate(
                    f"{median:.2f}",
                    (x, median),
                    xytext=(0, -17 if index else 12),
                    textcoords="offset points",
                    ha="center",
                    color=colors[index],
                    fontsize=10,
                    fontweight="bold",
                )
        before = comparison["baseline_summary"]["geomean_seconds"]
        after = comparison["candidate_summary"]["geomean_seconds"]
        ax.set_title(
            f"{title}\n{before:.2f} s → {after:.2f} s  (−{comparison['latency_reduction_pct']:.2f}%)",
            loc="left",
            fontsize=13,
            fontweight="bold",
            pad=18,
        )
        ax.set_xticks([0, 1], ["Product · seed 731", "Portrait · seed 1234"])
        ax.set_xlim(-0.48, 1.48)
        ax.set_ylim(58.5, 69)
        ax.grid(axis="y", color="#dce5ed", lw=0.8, zorder=0)
        ax.legend(frameon=False, loc="lower left", ncols=2)
        ax.spines["left"].set_color("#b7c8d6")
        ax.spines["bottom"].set_color("#b7c8d6")
    axes[0].set_ylabel("Warm complete generation · seconds")
    fig.suptitle(
        "2048 × 2048 · six steps · M5 Pro / 64 GiB · MLX 0.32.4",
        fontsize=14,
        fontweight="bold",
        y=0.98,
    )
    fig.text(
        0.5,
        0.025,
        "Dots: all three repeats. Lines: medians. Cached conditioning and resident weights; separate paired campaigns.\nHeadline values: geometric mean of case medians. Startup/downloads excluded; trace and PNG save included.",
        ha="center",
        fontsize=9,
        color="#465d70",
    )
    fig.tight_layout(rect=(0, 0.12, 1, 0.9), w_pad=3)
    target = root / "assets/benchmarks.svg"
    fig.savefig(target, facecolor=fig.get_facecolor())
    fig.savefig(
        "/tmp/minimax-benchmarks-preview.png", dpi=140, facecolor=fig.get_facecolor()
    )
    print(target)


if __name__ == "__main__":
    main()