#!/usr/bin/env python3
"""Nachmessung der im OMEGA-Hauptlauf abgeschnittenen Aufgaben.

Im Hauptlauf sind 10 von 150 Aufgaben ins 8000-Token-Limit gelaufen und haben
gar keine Antwort abgegeben. Damit ist offen, ob das Modell sie koennte oder
ob es sich festdenkt. Genau das misst dieses Skript: dieselben Aufgaben, volle
Denktiefe, ein Vielfaches an Budget.

Laeuft absichtlich gegen die PRODUKTION (PropellerA :8210). Das ist derselbe
Gewichtssatz (qwen3.8-27b-Q8_0.gguf), belegt einen von zwei Slots und braucht
keinen Stack-Abbau — der hat uns am 17.08. schon vier Stunden gekostet. Preis
dafuer: kein MTP-Drafter, also rund 20 statt 55 Token/s.

Zusaetzlich wird das Ende des Denktexts mitgeschrieben und auf Wiederholung
geprueft. Ein Modell, das im Kreis laeuft, ist etwas anderes als eines, dem
schlicht das Budget fehlt — und nur das erste waere das vielbeschriebene
"overthinking".
"""
import argparse
import json
import os
import sys
import time
import urllib.request

sys.path.insert(0, "/mnt/models/qwen38-prep")
import bench_mathe as M

QUELLE = "/mnt/models/qwen38-prep/omega/neubewertet/omega-reasoning.json"
ZIEL = "/mnt/models/qwen38-prep/omega/nachmessung.json"


def dreht_sich_im_kreis(text, fenster=300):
    """Grobe Schleifenerkennung: taucht ein spaeter Block frueher schon auf?

    Kein Beweis, aber ein brauchbarer Hinweis. Wir nehmen einen Block vom Ende
    und suchen ihn im vorderen Teil. Findet er sich, hat das Modell denselben
    Gedanken zweimal gedacht.
    """
    if len(text) < fenster * 3:
        return False
    block = text[-fenster:]
    return block in text[:-fenster]


def frage(port, modell, prompt, max_tokens, timeout):
    b = {"model": modell,
         "messages": [{"role": "user", "content": prompt}],
         "max_tokens": max_tokens, "temperature": 0.2, "top_p": 0.9,
         "chat_template_kwargs": {"reasoning_effort": "xhigh"}}
    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(tokens=0, sekunden=round(time.time() - t0, 1),
                            abbruch=f"fehler: {type(e).__name__}")
    m = d["choices"][0]["message"]
    t = d.get("timings", {})
    sek = time.time() - t0
    tok = t.get("predicted_n") or d.get("usage", {}).get("completion_tokens") or 0
    return (m.get("content") or ""), (m.get("reasoning_content") or ""), dict(
        tokens=tok, sekunden=round(sek, 1),
        tps=round(t.get("predicted_per_second") or (tok / sek if sek else 0), 1),
        abbruch=d["choices"][0].get("finish_reason") or "")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=8210)
    ap.add_argument("--modell", default="qwen3.8-27b")
    ap.add_argument("--max-tokens", type=int, default=32000)
    ap.add_argument("--timeout", type=int, default=3000)
    args = ap.parse_args()

    alt = json.load(open(QUELLE))
    offen = [x["id"] for x in alt["aufgaben"] if x.get("abbruch") == "length"]
    print(f"{len(offen)} abgeschnittene Aufgaben, Budget {args.max_tokens} "
          f"(vorher 8000), reasoning_effort=xhigh\n", flush=True)

    # Denselben Ausschnitt wie der Hauptlauf laden, sonst passen die IDs nicht.
    alle = {a["unique_id"] if "unique_id" in a else a.get("id"): a
            for a in M.lade("math500", 150)}
    erwartet_von = {x["id"]: x["erwartet"] for x in alt["aufgaben"]}

    ergebnisse = []
    geloest = 0
    for i, tid in enumerate(offen, 1):
        a = alle.get(tid)
        if a is None:
            print(f"[{i}/{len(offen)}] {tid}: nicht im Datensatz gefunden", flush=True)
            continue
        prompt = M.bau_prompt("math500", a)
        txt, denk, info = frage(args.port, args.modell, prompt,
                                args.max_tokens, args.timeout)
        gegeben = M.hole_boxed(txt)
        erwartet = erwartet_von[tid]
        ok = M.stimmt(gegeben, erwartet)
        geloest += 1 if ok else 0
        schleife = dreht_sich_im_kreis(denk)
        ergebnisse.append(dict(
            id=tid, ok=ok, gegeben=(gegeben or "")[:400], erwartet=erwartet,
            level=a.get("level"), schleife=schleife,
            denk_zeichen=len(denk), denk_ende=denk[-1200:], **info))
        print(f"[{i}/{len(offen)}] L{a.get('level')} "
              f"{'RICHTIG' if ok else 'weiterhin offen'}  "
              f"{info['tokens']} Tok  {info['abbruch']}  "
              f"{info['sekunden']}s{'  SCHLEIFE' if schleife else ''}  "
              f"-> {gegeben!r} (erwartet {erwartet!r})", flush=True)
        # Nach jeder Aufgabe schreiben: bei 10 langen Laeufen will man
        # Zwischenstaende sehen und nicht auf das Ende warten.
        json.dump(dict(quelle=QUELLE, max_tokens=args.max_tokens,
                       effort="xhigh", n=len(offen), geloest=geloest,
                       aufgaben=ergebnisse),
                  open(ZIEL, "w"), ensure_ascii=False, indent=1)

    am_limit = sum(1 for x in ergebnisse if x["abbruch"] == "length")
    schleifen = sum(1 for x in ergebnisse if x["schleife"])
    print(f"\n{geloest} von {len(ergebnisse)} jetzt geloest, "
          f"{am_limit} weiterhin am Limit, {schleifen} mit erkannter Schleife")
    print(f"geschrieben nach {ZIEL}")


if __name__ == "__main__":
    main()