#!/usr/bin/env python3
"""Run the standard consensus rule on formulas whose TeX is known and show every case for review.

    python3 tools/regeltest.py EPUB [--n 50] [--name regeltest] [--reuse vergleich/regeltest/ergebnis.json]
                                    [--exclude vergleich/regeltest/ergebnis.json ...]

Like `run`: pix2text reads every formula first, then Qwen (local); Gemma (local2) only where they
disagree; then the rough image check on pix2text's reading and decide_pair_tiebreak. Writes
vergleich/<name>/index.html (not versioned): summary, then per formula the original, the true
formula and every reading set with pdflatex, the decision and whether the result is right.
--reuse takes the readings (and their times) of an earlier run instead of asking the models again;
--exclude leaves out the formulas of earlier runs, so a new sample has none of them.
"""

from __future__ import annotations

import argparse
import html
import json
import random
import statistics
import sys
import time
from pathlib import Path

import numpy as np
from PIL import Image

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from konsens_bericht import orig_img, png_uri  # noqa: E402

from formula_img2latex.calibrate import _best_preset, corpus  # noqa: E402
from formula_img2latex.checks import pandoc_mathml_ok  # noqa: E402
from formula_img2latex.consensus import same_reading  # noqa: E402
from formula_img2latex.decide import consensus_thresholds, decide_pair_tiebreak  # noqa: E402
from formula_img2latex.engines import make_engine  # noqa: E402
from formula_img2latex.engines.p2t import tidy_p2t  # noqa: E402
from formula_img2latex.imageio import load_ink  # noqa: E402
from formula_img2latex.render import render_many  # noqa: E402
from formula_img2latex.tex import strip_delimiters  # noqa: E402
from formula_img2latex.verify import image_arbiter, verify  # noqa: E402

TOOLS = [("local", "Qwen"), ("p2t", "pix2text"), ("local2", "Gemma")]


def agreements(texes: list[str | None], display: bool, preset: str) -> list[tuple[int, int, str]]:
    """Pairs (Qwen, pix2text), (pix2text, Gemma), (Qwen, Gemma) as in pipeline.decide_all_consensus (strict)."""
    out = []
    for i, j in ((0, 1), (1, 2), (0, 2)):
        a, b = texes[i], texes[j]
        if a is None or b is None:
            continue
        how = same_reading(a, b, display, preset)
        if how:
            out.append((i, j, how))
    return out


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("epub", type=Path)
    ap.add_argument("--n", type=int, default=50)
    ap.add_argument("--name", default="regeltest")
    ap.add_argument("--reuse", type=Path, default=None, help="Lesarten aus einem früheren ergebnis.json")
    ap.add_argument("--refresh", default="",
                    help="diese Engines trotz --reuse neu fragen, z. B. local,local2 (neuer Prompt)")
    ap.add_argument("--arbiter", type=float, default=None,
                    help="Bild-Schiedsrichter einschalten mit dieser Schwelle (Qwen = Gemma gegen pix2text)")
    ap.add_argument("--exclude", type=Path, action="append", default=[],
                    help="Formeln aus diesem früheren ergebnis.json auslassen (mehrfach möglich)")
    args = ap.parse_args()
    cached = {}
    if args.reuse:
        cached = {c["name"]: c["readings"] for c in json.loads(args.reuse.read_text(encoding="utf-8"))["cases"]}
        for rs in cached.values():   # earlier runs stored pix2text's raw spaced output
            if rs[1] and rs[1].get("tex"):
                rs[1]["tex"] = tidy_p2t(rs[1].get("raw") or rs[1]["tex"])
    excluded = {c["name"] for x in args.exclude for c in json.loads(x.read_text(encoding="utf-8"))["cases"]}

    items = [it for it in corpus(args.epub) if it[0] not in excluded]
    smp = random.Random(42).sample(items, min(args.n + 1, len(items)))
    formulas, warm = smp[:args.n], smp[args.n]
    preset = _best_preset(smp[:12])
    th = consensus_thresholds({"arbiter": 1, "arbiter_margin": args.arbiter} if args.arbiter is not None else None)
    refresh = {e for e in args.refresh.split(",") if e}
    eng = {}
    for key, label in ([t for t in TOOLS if t[0] in refresh] if cached else TOOLS):
        e = make_engine(key)
        t0 = time.time()
        e.check()
        e.recognize(warm[1], warm[0], "", warm[3])   # warm-up, not counted
        eng[key] = e
        print(f"{label}: {e.model}, bereit in {time.time() - t0:.1f} s", file=sys.stderr, flush=True)

    cases = []
    for k, (name, data, truth, disp) in enumerate(formulas, 1):
        old = cached.get(name)
        if cached and old is None:
            raise SystemExit(f"{name} fehlt in {args.reuse} (andere Stichprobe?)")
        reads: list[dict | None] = [None, None, None]
        def read(pos: int) -> dict | None:
            key = TOOLS[pos][0]
            if old and key not in refresh:
                return old[pos]
            return eng[key].recognize(data, name, "", disp).to_json()

        for pos in (1, 0):   # pix2text first, then Qwen
            reads[pos] = read(pos)
        texes = [r and r.get("tex") for r in reads]
        agree = agreements(texes, disp, preset)
        if not any((i, j) == (0, 1) for i, j, _ in agree):   # disputed: ask Gemma
            reads[2] = read(2)
            texes = [r and r.get("tex") for r in reads]
            agree = agreements(texes, disp, preset)
        arbiter = None   # Qwen = Gemma against pix2text: image evidence for Qwen's reading
        pairs = {(i, j) for i, j, _ in agree}
        if (0, 2) in pairs and not pairs & {(0, 1), (1, 2)} and texes[0] and texes[1]:
            arbiter = image_arbiter(load_ink(data, name), texes[0], texes[1], disp, preset)
        _st, _why, use = decide_pair_tiebreak(reads, agree, None, True, None, th, arbiter)
        veto = None
        pandoc = True
        if use is not None and texes[use]:
            v = verify(load_ink(data, name), texes[use], disp, preset, check_variants=False)
            veto = {"ok": v.ok, "reason": v.reason, "metrics": v.metrics}
            pandoc = pandoc_mathml_ok([(texes[use], disp)])[0]
        status, reason, use = decide_pair_tiebreak(reads, agree, veto, pandoc, None, th, arbiter)
        ok_how = [same_reading(t, truth, disp, preset) if t else None for t in texes]
        ok = [bool(h) for h in ok_how]
        accepted = status == "auto_accepted"
        cases.append({"name": name, "truth": truth, "display": disp, "readings": reads, "ok": ok, "ok_how": ok_how,
                      "agree": agree, "arbiter": arbiter,
                      "status": status, "reason": reason, "use": use, "veto": veto,
                      "right": bool(accepted and ok[use])})
        mark = ("✓" if ok[use] else "✗ FALSCH") if accepted else "Review"
        print(f"{k:3d} {status:13s} {mark:8s} {reason[:110]}", file=sys.stderr, flush=True)

    n = len(cases)
    acc = [c for c in cases if c["status"] == "auto_accepted"]
    wrong = [c for c in acc if not c["right"]]
    review = [c for c in cases if c["status"] == "needs_review"]
    gemma = sum(c["readings"][2] is not None for c in cases)
    secs = {key: [c["readings"][p]["seconds"] for c in cases if c["readings"][p]] for p, (key, _l) in enumerate(TOOLS)}
    per_formula = [sum(r["seconds"] for r in c["readings"] if r) for c in cases]
    summary = {"formulas": n, "accepted": len(acc), "wrong": len(wrong), "review": len(review),
               "other": n - len(acc) - len(review), "gemma_asked": gemma,
               "mean_s_per_formula": statistics.mean(per_formula),
               "mean_s": {k: statistics.mean(v) if v else None for k, v in secs.items()},
               "correct_alone": {key: sum(c["ok"][p] for c in cases if c["readings"][p])
                                 for p, (key, _l) in enumerate(TOOLS)},
               # matches found only by the structural comparison (check by hand: the truth comparison uses it too)
               "structural": {"agree": sum(how == "strukturell" for c in cases for _i, _j, how in c["agree"]),
                              "truth": sum(h == "strukturell" for c in cases for h in c["ok_how"])}}

    out = Path(__file__).resolve().parents[1] / "vergleich" / args.name
    out.mkdir(parents=True, exist_ok=True)
    (out / "ergebnis.json").write_text(json.dumps({"epub": args.epub.name, "preset": preset, "summary": summary,
                                                   "cases": cases}, ensure_ascii=False, indent=1, default=str),
                                       encoding="utf-8")

    def order(c: dict) -> int:
        return 0 if c["status"] == "auto_accepted" and not c["right"] else 1 if c["status"] != "auto_accepted" else 2

    shown = sorted(cases, key=order)
    texes = []
    for c in shown:
        texes.append((strip_delimiters(c["truth"]), c["display"]))
        texes += [((r or {}).get("tex") or "\\text{–}", c["display"]) for r in c["readings"]]
    rendered = render_many(texes, preset)
    by_name = {f[0]: f[1] for f in formulas}
    lab = dict(TOOLS)
    head = (f"<h1>Regeltest: {html.escape(args.epub.name[:70])}</h1>"
            "<p>Regel: pix2text liest zuerst, Qwen bestätigt → annehmen. Sonst fragt Gemma; Gemma = pix2text → annehmen; "
            "Qwen = Gemma gegen pix2text → Mensch. Übernommen wird immer pix2texts Lesart (Leerzeichen bereinigt). "
            "Annahme nur mit Pandoc und Bild-Veto. "
            "✓/✗ = gleiche Zeichen an gleicher Stelle wie die wahre Formel (Abstände, kursiv/aufrecht, Fettdruck egal).</p>"
            "<table class=sum>"
            f"<tr><th>Formeln</th><td>{n}</td></tr>"
            f"<tr><th>automatisch angenommen</th><td>{len(acc)} ({len(acc) / n:.0%})</td></tr>"
            f"<tr><th>davon FALSCH</th><td><b>{len(wrong)}</b></td></tr>"
            f"<tr><th>Review (Mensch)</th><td>{len(review)} ({len(review) / n:.0%})</td></tr>"
            f"<tr><th>Gemma gefragt</th><td>{gemma} ({gemma / n:.0%})</td></tr>"
            f"<tr><th>nur strukturell gleich (Paare / gegen Wahrheit)</th><td>{summary['structural']['agree']} / "
            f"{summary['structural']['truth']}</td></tr>"
            f"<tr><th>Ø Zeit pro Formel (alle Tools nacheinander)</th><td>{summary['mean_s_per_formula']:.2f} s</td></tr>"
            + "".join(f"<tr><th>Ø {lab[k]}</th><td>{v:.2f} s</td></tr>" for k, v in summary["mean_s"].items() if v)
            + "".join(f"<tr><th>{lab[k]} allein richtig</th><td>{v}/{len(secs[k])}</td></tr>"
                      for k, v in summary["correct_alone"].items())
            + "</table>")
    parts = [head]
    per = 4
    for i, c in enumerate(shown):
        imgs = []
        for rr in rendered[per * i: per * (i + 1)]:
            if rr.ok:
                g = (rr.rasterize(2.0) * 255).astype(np.uint8)
                imgs.append(f"<img src='{png_uri(Image.fromarray(255 - g))}'>")
            else:
                imgs.append("<i>nicht setzbar</i>")
        if c["status"] == "auto_accepted":
            cls, verdict = ("good", "angenommen, richtig") if c["right"] else ("bad", "angenommen, FALSCH")
        elif c["status"] == "needs_review":
            cls, verdict = "rev", "Review"
        else:
            cls, verdict = "rev", c["status"]
        ncc = (c["veto"] or {}).get("metrics", {}).get("ncc")
        rows = [f"<tr><th>Original</th><td><img src='{png_uri(orig_img(by_name[c['name']], 2))}'></td><td></td></tr>",
                f"<tr><th>richtig</th><td>{imgs[0]}</td><td><code>{html.escape(c['truth'])}</code></td></tr>"]
        for p, (key, label) in enumerate(TOOLS):
            r = c["readings"][p]
            if r is None:
                rows.append(f"<tr><th>{label}</th><td><i>nicht gefragt</i></td><td></td></tr>")
                continue
            mark = "<b class=ok>✓</b>" if c["ok"][p] else "<b class=no>✗</b>"
            used = " <b>← verwendet</b>" if c["use"] == p and c["status"] == "auto_accepted" else ""
            rows.append(f"<tr><th>{mark} {label}{used}<br><small>{r['seconds']:.2f} s</small></th><td>{imgs[p + 1]}</td>"
                        f"<td><code>{html.escape(r.get('tex') or r.get('error') or '–')}</code>"
                        + (f"<br><small>roh: <code>{html.escape(r.get('raw') or '')}</code></small>"
                           if key == "p2t" and r.get("raw") and r.get("raw") != r.get("tex") else "")
                        + "</td></tr>")
        parts.append(f"<section class={cls}><div class=v>{verdict} · {html.escape(c['reason'])}"
                     f"{f' · Bild-Ähnlichkeit {ncc:.2f}' if ncc is not None else ''}</div>"
                     f"<small>{html.escape(c['name'])}</small><table>{''.join(rows)}</table></section>")
    (out / "index.html").write_text(
        "<!doctype html><meta charset=utf-8><title>Regeltest</title><style>"
        "body{font:14px system-ui,sans-serif;margin:16px;background:#f6f6f3}section{background:#fff;border:1px solid #ddd;"
        "border-left:6px solid #999;border-radius:6px;padding:6px 10px;margin:10px 0}section.good{border-left-color:#1f7a3a}"
        "section.bad{border-left-color:#a32020;background:#fff5f5}section.rev{border-left-color:#c88a00}"
        ".v{font-weight:600;font-size:13px}table.sum{border-collapse:collapse;background:#fff}"
        "table.sum td,table.sum th{border:1px solid #ccc;padding:4px 10px}"
        "th{text-align:left;font-weight:400;color:#555;width:190px;vertical-align:middle}td{padding:3px 8px;vertical-align:middle}"
        ".ok{color:#1f7a3a}.no{color:#a32020}code{font-size:12px}img{max-width:640px}</style>" + "".join(parts),
        encoding="utf-8")
    print(json.dumps(summary, ensure_ascii=False, indent=1))
    print(out / "index.html")
    return 0


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