#!/usr/bin/env python3
"""Render primary KMeans latency panels from the preserved measurement JSON.

Uses Matplotlib when installed, with a dependency-free SVG renderer otherwise.
Both renderers use medians/MAD recomputed from raw seconds, zero-origin axes,
and distinct legacy/new-serial/reused-pool rows. Never runs a benchmark.
"""

from __future__ import annotations

import argparse
import hashlib
from html import escape
import importlib.util
import json
import math
from pathlib import Path
import statistics

MODES = [("legacy", "Legacy serial", "#8b97a3"),
         ("serial", "New serial", "#173b4f"),
         (1, "Pool: 1 worker", "#478c9a"),
         (2, "Pool: 2 workers", "#287758"),
         (4, "Pool: 4 workers", "#b47521")]


def observations(workload: dict, operation: str) -> list[dict]:
    study = workload[operation]
    pools = {row["workers"]: row for row in study["pools"]}
    if set(pools) != {1, 2, 4}:
        raise ValueError("expected exactly the reused 1/2/4-worker modes")
    result = []
    for mode, label, color in MODES:
        row = study[mode] if isinstance(mode, str) else pools[mode]
        samples = row["raw_seconds"]
        if len(samples) != 7 or any(not math.isfinite(v) or v <= 0 for v in samples):
            raise ValueError("expected seven finite positive retained observations")
        center = statistics.median(samples)
        spread = statistics.median(abs(value - center) for value in samples)
        for actual, expected in [(center, row["median_seconds"]), (spread, row["mad_seconds"])]:
            if not math.isclose(actual, expected, rel_tol=1e-12, abs_tol=1e-15):
                raise ValueError("recorded summary differs from raw observations")
        result.append({"label": label, "color": color, "ms": center * 1000,
                       "mad_ms": spread * 1000})
    return result


def svg(primary: list[dict], source: Path, digest: str) -> str:
    parts = ['<svg xmlns="http://www.w3.org/2000/svg" width="1440" height="760" '
             'viewBox="0 0 1440 760" role="img" aria-labelledby="title description">',
             '<title id="title">KMeans primary fit and prediction latency</title>',
             '<desc id="description">Batch and Wide workloads: legacy serial, updated serial, '
             'and reused pools with 1, 2 and 4 workers. Bars begin at zero. Values and whiskers '
             'show median plus or minus median absolute deviation in milliseconds, not confidence '
             'intervals. Complete data and other workloads appear in the accompanying report.</desc>',
             '<metadata>' + escape(json.dumps({"source": source.name, "sha256": digest,
                 "renderer": "tools/plot_kmeans_benchmark.py:svg",
                 "units": "milliseconds", "spread": "median absolute deviation"}, sort_keys=True)) + '</metadata>',
             '<rect width="1440" height="760" fill="#ffffff"/>',
             '<g font-family="Arial, Helvetica, sans-serif" fill="#18352c">']

    def text(x: float, y: float, value: str, size: int = 13, anchor: str = "start", weight: int = 400):
        parts.append(f'<text x="{x:.2f}" y="{y:.2f}" font-size="{size}" '
                     f'text-anchor="{anchor}" font-weight="{weight}">{escape(value)}</text>')

    text(32, 42, "KMeans: recorded local CPU latency", 27, weight=600)
    text(32, 69, "Apple M2 Max · macOS arm64 · Release/ThinLTO · 7 retained observations after 2 warmups", 14)
    for panel, operation in enumerate(("fit", "predict")):
        left = 32 + panel * 704
        origin, width = left + 120, 400
        values = [observations(workload, operation) for workload in primary]
        ceiling = max(row["ms"] + row["mad_ms"] for group in values for row in group)
        magnitude = 10 ** math.floor(math.log10(ceiling))
        step = magnitude / 5
        maximum = math.ceil(ceiling / step) * step
        text(left, 118, "Fitting" if operation == "fit" else "Prediction", 22, weight=600)
        for index in range(6):
            value = maximum * index / 5
            x = origin + value / maximum * width
            parts.append(f'<line x1="{x:.2f}" x2="{x:.2f}" y1="164" y2="519" '
                         'stroke="#dce5e0" stroke-width="1"/>')
            text(x, 549, f"{value:g}", 12, "middle")
        text(origin + width / 2, 576, "Median latency (ms); lower is faster", 13, "middle")
        for group, (workload, rows) in enumerate(zip(primary, values)):
            heading = 151 + group * 196
            text(left, heading, f'{workload["name"]}  ·  N={workload["n"]:,}, '
                 f'Q={workload["q"]:,}, D={workload["d"]}, k={workload["k"]}', 12, weight=600)
            for index, row in enumerate(rows):
                y = heading + 32 + index * 29
                end = origin + row["ms"] / maximum * width
                lower = origin + max(0, row["ms"] - row["mad_ms"]) / maximum * width
                upper = origin + (row["ms"] + row["mad_ms"]) / maximum * width
                text(origin - 10, y + 4, row["label"], 12, "end")
                parts.append(f'<rect x="{origin:.2f}" y="{y-8:.2f}" width="{end-origin:.2f}" '
                             f'height="16" fill="{row["color"]}"/>')
                parts.append(f'<path d="M {lower:.2f} {y:.2f} H {upper:.2f} '
                             f'M {lower:.2f} {y-5:.2f} V {y+5:.2f} '
                             f'M {upper:.2f} {y-5:.2f} V {y+5:.2f}" '
                             'stroke="#10261d" fill="none" stroke-width="1.2"/>')
                text(end + 8, y + 4, f'{row["ms"]:.3f} ± {row["mad_ms"]:.3f}', 12)
    text(32, 619, "Bars and labels: median ± MAD in ms. MAD describes observed spread; it is not a confidence interval.", 14)
    text(32, 646, "Pools are reused. Serial creates no pool. Legacy and updated serial use different numerical accumulation algorithms.", 14)
    text(32, 673, "Development runtime observations on one machine; these measurements establish neither universal nor GPU/cluster speedups.", 14)
    text(32, 711, "Source: " + source.name + " · Complete values, tiny/cold cases, workspace and process RSS remain in the report/JSON.", 12)
    parts.append('</g></svg>\n')
    return '\n'.join(parts)


def matplotlib(primary: list[dict], output: Path, source: Path, digest: str) -> None:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    plt.rcParams.update({"font.family": "sans-serif", "font.size": 11, "svg.fonttype": "none"})
    figure, axes = plt.subplots(1, 2, figsize=(16, 8))
    for axis, operation in zip(axes, ("fit", "predict")):
        ticks, labels = [], []
        for group, workload in enumerate(primary):
            rows = observations(workload, operation)
            for index, row in enumerate(rows):
                y = group * 7 + index
                axis.barh(y, row["ms"], xerr=row["mad_ms"], color=row["color"],
                          height=.65, capsize=3, error_kw={"elinewidth": 1})
                axis.annotate(f'{row["ms"]:.3f} ± {row["mad_ms"]:.3f}',
                              (row["ms"] + row["mad_ms"], y), xytext=(6, 0),
                              textcoords="offset points", va="center", fontsize=9)
                ticks.append(y)
                labels.append(workload["name"] + " · " + row["label"])
        axis.set_yticks(ticks, labels)
        axis.invert_yaxis()
        axis.set_xlim(left=0)
        axis.margins(x=.27)
        axis.set_xlabel("Median latency (ms); lower is faster")
        axis.set_title("Fitting" if operation == "fit" else "Prediction", loc="left", fontsize=17)
        axis.set_axisbelow(True)
        axis.grid(axis="x", alpha=.25)
        for edge in ("top", "right"):
            axis.spines[edge].set_visible(False)
    figure.suptitle("KMeans: recorded local CPU latency", x=.03, ha="left", fontsize=22)
    figure.text(.03, .92, "Apple M2 Max · macOS arm64 · Release/ThinLTO · 7 retained observations after 2 warmups")
    figure.text(.03, .055, "Median ± MAD, not confidence intervals. Reused pools; legacy/new serial arithmetic differs. Development runtime, one recorded machine.")
    figure.tight_layout(rect=(0, .09, 1, .89), w_pad=3)
    figure.savefig(output, format="svg", metadata={"Title": "KMeans primary fit and prediction latency",
        "Description": "Source " + source.name + "; SHA-256 " + digest,
        "Creator": "tools/plot_kmeans_benchmark.py", "Date": None})
    plt.close(figure)


def main() -> None:
    root = Path(__file__).resolve().parents[1]
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", type=Path, default=root / "docs/benchmarks/kmeans_native_20261030.json")
    parser.add_argument("--output", type=Path, default=root / "docs/benchmarks/kmeans_native_20261030.svg")
    parser.add_argument("--renderer", choices=("auto", "matplotlib", "svg"), default="auto")
    args = parser.parse_args()
    raw = args.input.read_bytes()
    study = json.loads(raw)
    if study.get("schema") != "aner.kmeans-evidence/1" or study.get("status") != "verified":
        parser.error("expected verified aner.kmeans-evidence/1 input")
    primary = [next(workload for workload in study["workloads"] if workload["id"] == name)
               for name in ("batch", "wide")]
    digest = hashlib.sha256(raw).hexdigest()
    available = importlib.util.find_spec("matplotlib") is not None
    backend = args.renderer if args.renderer != "auto" else "matplotlib" if available else "svg"
    if backend == "matplotlib" and not available:
        parser.error("Matplotlib is unavailable; use --renderer svg without installing dependencies")
    args.output.parent.mkdir(parents=True, exist_ok=True)
    if backend == "svg":
        args.output.write_text(svg(primary, args.input, digest))
    else:
        matplotlib(primary, args.output, args.input, digest)
    print(json.dumps({"renderer": backend, "source_sha256": digest, "output": str(args.output)}))


if __name__ == "__main__":
    main()