#!/usr/bin/env python3
"""Tracer les deux séries réelles avec matplotlib (stock.json, unlocked.json)."""

import argparse
import json
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Patch


parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("stock", type=Path)
parser.add_argument("unlocked", type=Path)
parser.add_argument("output", type=Path)
args = parser.parse_args()
datasets = [json.loads(path.read_text()) for path in (args.stock, args.unlocked)]
profiles = ("short", "medium", "long")
assert all(d["status"] == "complete" and len(d["runs"]) == 9 for d in datasets)
for before, after in zip(datasets[0]["runs"], datasets[1]["runs"]):
    assert (before["profile"], before["repeat"], before["prompt_sha256"]) == (
        after["profile"], after["repeat"], after["prompt_sha256"]
    )

plt.rcParams.update({
    "font.family": "DejaVu Sans", "font.size": 14,
    "axes.spines.top": False, "axes.spines.right": False,
})
fig, axes = plt.subplots(1, 2, figsize=(16, 9), dpi=100)
background, foreground = "#f7f8fa", "#182230"
colors = ("#64748b", "#087f8c")
fig.patch.set_facecolor(background)
fig.subplots_adjust(top=.69, bottom=.23, left=.075, right=.96, wspace=.25)
fig.text(.075, .92, "gpu-unlock sur RTX 3090", fontsize=29, weight="bold", color=foreground)
fig.text(.075, .865, "Qwen3.8 27B Q4_K_M avec MTP : avant et après", fontsize=20, color="#475467")
fig.text(.075, .815, "Ollama 0.33.2 · pilote 550.163.01 · limite de puissance 350 W", fontsize=15, color="#475467")
fig.legend(
    [Patch(color=c) for c in colors], ["Sans gpu-unlock", "Avec gpu-unlock"],
    loc="upper left", bbox_to_anchor=(.068, .789), ncol=2, frameon=False, fontsize=14,
)
labels = [str(datasets[0]["summary"][p]["prompt_tokens"][0]) for p in profiles]
labels = [f"{int(value):,}".replace(",", " ") for value in labels]
for ax, key, title in zip(axes, ("prompt_tokens_per_second", "generation_tokens_per_second"),
                          ("Traitement du prompt", "Génération de la réponse")):
    ax.set_facecolor(background)
    top = max(d["summary"][p][key]["max"] for d in datasets for p in profiles)
    ax.set_ylim(0, top * 1.3)
    for index, (data, color) in enumerate(zip(datasets, colors)):
        values = [data["summary"][p][key]["median"] for p in profiles]
        low = [v - data["summary"][p][key]["min"] for p, v in zip(profiles, values)]
        high = [data["summary"][p][key]["max"] - v for p, v in zip(profiles, values)]
        positions = [n + (-.18 if index == 0 else .18) for n in range(3)]
        ax.bar(positions, values, width=.32, color=color, yerr=[low, high], capsize=4,
               error_kw={"elinewidth": 1.3, "ecolor": foreground}, zorder=3)
        for x, p, value in zip(positions, profiles, values):
            label = f"{value:.0f}" if key.startswith("prompt") else f"{value:.1f}".replace(".", ",")
            ax.annotate(label, (x, data["summary"][p][key]["max"]), xytext=(0, 8),
                        textcoords="offset points", ha="center", fontsize=13, weight="bold", color=color)
    for n, profile in enumerate(profiles):
        before = datasets[0]["summary"][profile][key]["median"]
        after = datasets[1]["summary"][profile][key]["median"]
        delta = (after / before - 1) * 100
        y = max(d["summary"][profile][key]["max"] for d in datasets) + top * .12
        ax.text(n, y, f"{delta:+.1f} %".replace(".", ","), ha="center", fontsize=14,
                weight="bold", color=foreground)
    ax.set_xticks(range(3), labels)
    ax.set_xlabel("Tokens en entrée", labelpad=12)
    ax.set_ylabel("Tokens par seconde", labelpad=10)
    ax.set_title(title, loc="left", weight="bold", fontsize=18, pad=18)
    ax.yaxis.grid(True, color="#dce1e8", zorder=0)
    ax.set_axisbelow(True)
    ax.tick_params(axis="both", length=0, pad=10)
fig.text(.075, .115, "Médianes de 3 essais par profil et par état ; barres : minimum et maximum.", fontsize=14, color="#475467")
fig.text(.075, .075, "Fréquences automatiques · contexte 16 384 · sortie de 256 tokens · MTP 4", fontsize=12, color="#475467")
fig.text(.96, .035, "cryptolab.re · 24 octobre 2026", ha="right", fontsize=12, color="#667085")
fig.savefig(args.output, dpi=100, facecolor=fig.get_facecolor())
plt.close(fig)
print(args.output)