#!/usr/bin/env python3
"""MATH-500 und GSM8K gegen einen OpenAI-kompatiblen Server.

Warum Mathe zusaetzlich zu HumanEval+: dort erreichten alle sieben Varianten
88-91 Prozent und waren statistisch ununterscheidbar. Mathematik trennt
staerker, vor allem bei den hoeheren Schwierigkeitsstufen, und ist der Ort,
an dem sich zeigt, ob ein Denkblock etwas bringt.

Die Bewertung ist die eigentliche Arbeit: MATH-500 erwartet LaTeX-Antworten
(`\\frac{3}{4}`, `\\left( 3, \\frac{\\pi}{2} \\right)`), und dieselbe Zahl kann
in einem Dutzend Schreibweisen kommen. Deshalb wird normalisiert statt
stur verglichen — und wo Normalisierung nicht reicht, zusaetzlich numerisch
ausgewertet.
"""
import argparse
import json
import os
import re
import sys
import time
import urllib.request
from fractions import Fraction

DS = "/mnt/models/qwen38-prep/datasets"
REASONING = False


def lade(satz, grenze):
    zeilen = [json.loads(z) for z in open(os.path.join(DS, f"{satz}.jsonl"))
              if z.strip()]
    if not grenze or grenze >= len(zeilen):
        return zeilen
    # Gleichmaessig ueber den Datensatz greifen statt die ersten N: MATH-500
    # ist nach Thema sortiert, die ersten 150 waeren fast nur Precalculus.
    schritt = len(zeilen) / grenze
    return [zeilen[int(i * schritt)] for i in range(grenze)]


def bau_prompt(satz, a):
    if satz.startswith("math"):
        return (a["problem"] + "\n\nLöse die Aufgabe. Schreibe die "
                "endgültige Antwort ganz zum Schluss in der Form "
                "\\boxed{ANTWORT}.")
    return (a["question"] + "\n\nLöse die Aufgabe Schritt für Schritt. "
            "Schreibe die endgültige Zahl ganz zum Schluss in der Form "
            "#### ZAHL.")


def frage(port, prompt, max_tokens, timeout):
    b = {"messages": [{"role": "user", "content": prompt}],
         "max_tokens": max_tokens, "cache_prompt": False,
         "temperature": 0.2, "top_p": 0.9}
    if not REASONING:
        b["chat_template_kwargs"] = {"enable_thinking": False}
    req = urllib.request.Request(
        f"http://127.0.0.1:{port}/v1/chat/completions",
        data=json.dumps(b).encode(),
        headers={"Content-Type": "application/json"})
    t0 = time.time()
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            d = json.load(r)
    except Exception as e:
        return "", dict(tps=0.0, tokens=0, sekunden=time.time() - t0,
                        abbruch="fehler"), f"{type(e).__name__}: {e}"
    m = d["choices"][0]["message"]
    t = d.get("timings", {})
    # vLLM liefert keine timings — dann aus usage und Wanduhr rechnen.
    sek = time.time() - t0
    tok = t.get("predicted_n") or d.get("usage", {}).get("completion_tokens") or 0
    tps = t.get("predicted_per_second") or (tok / sek if sek > 0 else 0.0)
    return (m.get("content") or ""), dict(
        tps=tps, tokens=tok,
        denk_zeichen=len(m.get("reasoning_content") or ""),
        abbruch=d["choices"][0].get("finish_reason") or "",
        sekunden=round(sek, 2)), ""


# --- Antwortextraktion -----------------------------------------------------

def hole_boxed(text):
    """Inhalt des LETZTEN \\boxed{...}. Klammern werden mitgezaehlt, weil
    verschachtelte Ausdruecke wie \\boxed{\\frac{1}{2}} sonst abbrechen."""
    stelle = text.rfind("\\boxed")
    if stelle < 0:
        return None
    i = text.find("{", stelle)
    if i < 0:
        return None
    tiefe, j = 0, i
    while j < len(text):
        if text[j] == "{":
            tiefe += 1
        elif text[j] == "}":
            tiefe -= 1
            if tiefe == 0:
                return text[i + 1:j]
        j += 1
    return None


def hole_gsm(text):
    m = re.findall(r"####\s*([-\d.,]+)", text)
    if m:
        return m[-1]
    # Kein Marker: die letzte Zahl im Text nehmen.
    z = re.findall(r"-?\d[\d.,]*", text)
    return z[-1] if z else None


def normalisiere(s):
    """Bringt LaTeX-Antworten auf eine vergleichbare Form.

    Die Faelle hier stammen aus echten Fehlschlaegen des ersten Durchgangs:
    von acht als falsch gewerteten Antworten war genau eine wirklich falsch,
    die uebrigen sieben waren Schreibweisen. `\\frac43` gegen `\\frac{4}{3}`,
    `.35625` gegen `0,35625`, `x=5` gegen `5`. Wer das nicht abfaengt, misst
    LaTeX-Konventionen statt Mathematik.
    """
    if s is None:
        return None
    s = s.strip()
    # \frac43 und \frac 34 -> \frac{4}{3}: die Kurzform ohne Klammern ist in
    # MATH-500 haeufig und sonst nicht mit der Langform vergleichbar.
    s = re.sub(r"\\(d|t)?frac\s*(\d)\s*(\d)", r"\\frac{\2}{\3}", s)
    for weg in ("\\left", "\\right", "\\!", "\\,", "\\;", "\\ ", "$", " "):
        s = s.replace(weg, "")
    s = s.replace("dfrac", "frac").replace("tfrac", "frac")
    s = s.replace("^{\\circ}", "").replace("^\\circ", "")
    s = re.sub(r"\\text\{([^}]*)\}", r"\1", s)
    s = re.sub(r"\\mbox\{([^}]*)\}", r"\1", s)
    # Zuweisungspraefixe: "x=5" und "5" sind dieselbe Antwort, ebenso
    # "x\in[-2,7]" und "[-2,7]". Nur bei einem einzelnen Buchstaben links,
    # damit Gleichungen als Antwort nicht zerstoert werden.
    s = re.sub(r"^[a-zA-Z]\s*(=|\\in)\s*", "", s)
    # Multiple-Choice: "(B)" und "B" sind dasselbe.
    m = re.fullmatch(r"\(([a-eA-E])\)", s)
    if m:
        s = m.group(1)
    # Komma in Zahlen: Tausendertrenner oder Dezimaltrenner? Unterschieden
    # wird an der Stellenzahl — "1,000" ist Tausender (englische Konvention
    # des Datensatzes), "0,35625" kann es nicht sein. Die Dezimalpruefung
    # MUSS zuerst kommen, sonst frisst die Tausenderregel das Komma weg.
    if re.fullmatch(r"-?\d*,\d+", s) and not re.fullmatch(r"-?\d+,\d{3}", s):
        s = s.replace(",", ".")
    elif re.fullmatch(r"-?[\d,]+(\.\d+)?", s):
        s = s.replace(",", "")
    # Fuehrende Null: ".5" und "0.5" sind dieselbe Zahl.
    s = re.sub(r"^(-?)\.(\d)", r"\g<1>0.\2", s)
    s = s.rstrip(".")
    # Einheiten und Prozentzeichen am Ende sind fuer die Gleichheit egal.
    s = re.sub(r"(\\%|%)$", "", s)
    return s.lower()


def als_zahl(s):
    """Versucht, einen Ausdruck numerisch auszuwerten — nur die Formen, die in
    diesen Datensaetzen vorkommen. Kein eval, das waere ein Sicherheitsloch."""
    if s is None:
        return None
    s = s.strip()
    m = re.fullmatch(r"\\frac\{(-?[\d.]+)\}\{(-?[\d.]+)\}", s)
    if m:
        try:
            return float(Fraction(m.group(1)) / Fraction(m.group(2)))
        except Exception:
            return None
    m = re.fullmatch(r"(-?[\d.]+)/(-?[\d.]+)", s)
    if m:
        try:
            return float(Fraction(m.group(1)) / Fraction(m.group(2)))
        except Exception:
            return None
    try:
        return float(s)
    except ValueError:
        return None


def stimmt(gegeben, erwartet):
    a, b = normalisiere(gegeben), normalisiere(erwartet)
    if a is None or b is None:
        return False
    if a == b:
        return True
    za, zb = als_zahl(a), als_zahl(b)
    if za is not None and zb is not None:
        return abs(za - zb) < 1e-6
    return False


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=8291)
    ap.add_argument("--modell", required=True)
    ap.add_argument("--satz", default="math500")
    ap.add_argument("--out", required=True)
    ap.add_argument("--grenze", type=int, default=0)
    ap.add_argument("--max-tokens", type=int, default=1600)
    ap.add_argument("--timeout", type=int, default=600)
    ap.add_argument("--reasoning", action="store_true")
    args = ap.parse_args()

    global REASONING
    REASONING = args.reasoning
    aufgaben = lade(args.satz, args.grenze)
    print(f"{args.modell} | {args.satz} | {len(aufgaben)} Aufgaben | "
          f"Reasoning {'AN' if REASONING else 'AUS'}", file=sys.stderr)

    zeilen, ok_n = [], 0
    t0 = time.time()
    for i, a in enumerate(aufgaben, 1):
        erwartet = a.get("answer") if args.satz.startswith("math") \
            else hole_gsm(a.get("answer", ""))
        txt, m, fehler = frage(args.port, bau_prompt(args.satz, a),
                               args.max_tokens, args.timeout)
        if fehler:
            zeilen.append(dict(id=str(a.get("unique_id", i)), ok=False,
                               grund=fehler, **m))
            print(f"  [{i:3d}] SERVERFEHLER {fehler[:50]}", file=sys.stderr)
            continue
        gegeben = hole_boxed(txt) if args.satz.startswith("math") \
            else hole_gsm(txt)
        ok = stimmt(gegeben, erwartet)
        ok_n += ok
        zeilen.append(dict(id=str(a.get("unique_id", i)), ok=ok,
                           level=a.get("level"), thema=a.get("subject"),
                           gegeben=(gegeben or "")[:60],
                           erwartet=(str(erwartet) or "")[:60], **m))
        if i % 20 == 0 or not ok:
            print(f"  [{i:3d}/{len(aufgaben)}] "
                  f"{'OK' if ok else 'FEHL: ' + repr(gegeben)[:26] + ' statt ' + repr(erwartet)[:26]}"
                  f"   (bisher {ok_n}/{i} = {ok_n/i*100:.0f}%)", file=sys.stderr)

    dauer = time.time() - t0
    tps = sorted(x["tps"] for x in zeilen if x["tps"] > 0)
    # Aufschluesselung nach Schwierigkeitsgrad — der eigentliche Mehrwert
    # gegenueber einer einzelnen Prozentzahl.
    nach_level = {}
    for x in zeilen:
        lv = x.get("level")
        if lv is not None:
            d = nach_level.setdefault(str(lv), [0, 0])
            d[1] += 1
            d[0] += 1 if x["ok"] else 0
    bericht = dict(
        modell=args.modell, satz=args.satz, reasoning=REASONING,
        n=len(aufgaben), bestanden=ok_n,
        quote=round(ok_n / len(aufgaben) * 100, 1) if aufgaben else 0.0,
        tps_median=round(tps[len(tps) // 2], 1) if tps else 0.0,
        tokens_je_aufgabe=round(sum(x["tokens"] for x in zeilen) / len(zeilen), 1)
                          if zeilen else 0,
        am_limit=sum(1 for x in zeilen if x.get("abbruch") == "length"),
        nach_level={k: v for k, v in sorted(nach_level.items())},
        dauer_s=round(dauer, 1), aufgaben=zeilen)
    with open(args.out, "w") as f:
        json.dump(bericht, f, ensure_ascii=False, indent=1)
    lv = "  ".join(f"L{k}: {v[0]}/{v[1]}" for k, v in bericht["nach_level"].items())
    print(f"\n{args.modell} | {args.satz}: {ok_n}/{len(aufgaben)} = "
          f"{bericht['quote']}%  |  {bericht['tps_median']} t/s  |  "
          f"{bericht['tokens_je_aufgabe']:.0f} Token/Aufgabe  |  "
          f"{bericht['am_limit']} abgeschnitten  |  {dauer/60:.1f} min",
          file=sys.stderr)
    if lv:
        print(f"  nach Schwierigkeit: {lv}", file=sys.stderr)


if __name__ == "__main__":
    main()