#!/usr/bin/env python3
"""Mathe-Diagramme fuer die HuggingFace-Seite.

Zwei Bilder:
  1. MATH-500 je Variante mit 95-Prozent-Intervall — zeigt, dass die
     Varianten sich nicht unterscheiden
  2. Trefferquote nach Schwierigkeitsstufe — zeigt, dass der Benchmark
     sehr wohl differenziert, nur eben nicht zwischen diesen Modellen

Bild 2 ist das eigentlich interessante: eine einzelne Prozentzahl verschweigt,
dass die Modelle auf Stufe 1 bei knapp 90 Prozent liegen und auf Stufe 5 bei
gut der Haelfte.
"""
import glob
import json
import math
import os
import sys

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

sys.path.insert(0, "/mnt/models/qwen38-prep")
from diagramm import HELL, DUNKEL, NAMEN, stil

PREP = "/mnt/models/qwen38-prep"


def lade():
    d = {}
    for p in glob.glob(os.path.join(PREP, "mathe-math500-neubewertet", "*.json")):
        k = os.path.splitext(os.path.basename(p))[0]
        if k in NAMEN or k in ("q36-35b-a3b", "gemma4-12b"):
            d[k] = json.load(open(p))
    return d


def kurzname(k, j):
    if k in NAMEN:
        return NAMEN[k][0]
    return j["modell"].split(" (")[0]


def ist38(k):
    return NAMEN[k][1] if k in NAMEN else False


def bild_quoten(d, f, ziel):
    zeilen = sorted(
        ((kurzname(k, j), ist38(k), j["quote"], j["bestanden"], j["n"],
          j.get("am_limit", 0))
         for k, j in d.items()), key=lambda z: z[2])
    if not zeilen:
        return False

    fig, ax = plt.subplots(figsize=(8.6, 0.55 * len(zeilen) + 2.0), dpi=170)
    fig.patch.set_facecolor(f["surface"])
    y = list(range(len(zeilen)))
    ax.barh(y, [z[2] for z in zeilen], height=0.58, zorder=3,
            color=[f["akzent"] if z[1] else f["grau"] for z in zeilen])
    hoechst = max(z[2] for z in zeilen)
    for i, z in enumerate(zeilen):
        p = z[3] / z[4]
        ci = 1.96 * math.sqrt(p * (1 - p) / z[4]) * 100
        ax.plot([z[2] - ci, z[2] + ci], [i, i], color=f["ink2"],
                linewidth=1.2, zorder=5)
        for x in (z[2] - ci, z[2] + ci):
            ax.plot([x, x], [i - 0.13, i + 0.13], color=f["ink2"],
                    linewidth=1.2, zorder=5)
        ax.text(z[2] + ci + hoechst * 0.02, i, f"{z[2]:.1f} %", va="center",
                ha="left", color=f["ink"], fontsize=10, fontweight="600")
        ax.text(hoechst * 0.025, i, f"{z[3]}/{z[4]}", va="center", ha="left",
                color=f["surface"], fontsize=8.5, zorder=6)
    ax.set_yticks(y)
    ax.set_yticklabels([z[0] for z in zeilen], color=f["ink"], fontsize=10)
    ax.set_xlim(0, hoechst * 1.30)
    stil(ax, f, "MATH-500  ·  gelöste Aufgaben",
         "Anteil gelöst in Prozent  ·  Strich = 95-Prozent-Intervall")
    abgeschnitten = sum(z[5] for z in zeilen) / len(zeilen)
    fig.text(0.012, 0.015,
             f"Alle Intervalle überlappen.  Im Schnitt {abgeschnitten:.0f} von "
             f"{zeilen[0][4]} Antworten liefen ins Token-Limit und zählen als "
             f"falsch — die absoluten Werte liegen daher zu niedrig.",
             color=f["muted"], fontsize=8)
    fig.tight_layout(rect=(0, 0.05, 1, 1))
    fig.savefig(ziel, facecolor=f["surface"])
    plt.close(fig)
    return True


def bild_stufen(d, f, ziel):
    """Trefferquote nach Schwierigkeitsstufe. Hier trennt der Benchmark."""
    stufen = sorted({s for j in d.values() for s in j.get("nach_level", {})})
    if not stufen:
        return False
    fig, ax = plt.subplots(figsize=(8.6, 5.2), dpi=170)
    fig.patch.set_facecolor(f["surface"])
    # Farbe traegt die Generation, Strichart und Markerform die Variante.
    # Farbe allein reicht hier nicht: mehrere Qwen3.8-Varianten teilen sich
    # denselben Blauton und waeren in einem Liniendiagramm nicht trennbar.
    striche = ["-", "--", ":", "-."]
    marker = ["o", "s", "^", "D"]
    zaehler = {True: 0, False: 0}
    for k, j in sorted(d.items()):
        nl = j.get("nach_level", {})
        xs, ys = [], []
        for s in stufen:
            if s in nl and nl[s][1]:
                xs.append(int(s))
                ys.append(nl[s][0] / nl[s][1] * 100)
        if not xs:
            continue
        g = ist38(k)
        i = zaehler[g] % 4
        zaehler[g] += 1
        ax.plot(xs, ys, marker=marker[i], markersize=7, linewidth=2,
                linestyle=striche[i],
                color=f["akzent"] if g else f["grau"],
                markeredgecolor=f["surface"], markeredgewidth=1.5,
                label=kurzname(k, j), zorder=3)
    ax.set_facecolor(f["surface"])
    ax.set_title("MATH-500 nach Schwierigkeitsstufe", color=f["ink"],
                 fontsize=12.5, pad=12, loc="left", fontweight="600")
    ax.set_xlabel("Schwierigkeitsstufe (1 = leicht, 5 = Wettbewerbsniveau)",
                  color=f["ink2"], fontsize=9.5)
    ax.set_ylabel("gelöst in Prozent", color=f["ink2"], fontsize=9.5)
    ax.set_xticks([int(s) for s in stufen])
    ax.tick_params(colors=f["muted"], labelsize=9.5, length=0)
    for s in ("top", "right"):
        ax.spines[s].set_visible(False)
    for s in ("left", "bottom"):
        ax.spines[s].set_color(f["achse"])
    ax.grid(True, color=f["grid"], linewidth=1)
    ax.set_axisbelow(True)
    ax.set_ylim(0, 100)
    leg = ax.legend(frameon=False, fontsize=9, loc="lower left")
    for t in leg.get_texts():
        t.set_color(f["ink2"])
    fig.text(0.012, 0.015,
             "Der Benchmark trennt sauber über die Stufen — nur eben nicht "
             "zwischen diesen Varianten.  Blau: Qwen3.8, Grau: übrige",
             color=f["muted"], fontsize=8)
    fig.tight_layout(rect=(0, 0.045, 1, 1))
    fig.savefig(ziel, facecolor=f["surface"])
    plt.close(fig)
    return True


def main():
    d = lade()
    print(f"MATH-500: {len(d)} Modelle")
    if not d:
        return 1
    ziel = os.path.join(PREP, "bilder")
    os.makedirs(ziel, exist_ok=True)
    for name, f in (("hell", HELL), ("dunkel", DUNKEL)):
        for bau, datei in ((bild_quoten, f"mathe-{name}.png"),
                           (bild_stufen, f"mathe-stufen-{name}.png")):
            if bau(d, f, os.path.join(ziel, datei)):
                print(f"  {datei}")
    return 0


if __name__ == "__main__":
    sys.exit(main())