#!/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 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("--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()
    eng = {}
    for key, label in ([] 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]
        for pos in (1, 0):   # pix2text first, then Qwen
            reads[pos] = old[pos] if old else eng[TOOLS[pos][0]].recognize(data, name, "", disp).to_json()
        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] = old[2] if old else eng["local2"].recognize(data, name, "", disp).to_json()
            texes = [r and r.get("tex") for r in reads]
            agree = agreements(texes, disp, preset)
        _st, _why, use = decide_pair_tiebreak(reads, agree, None, True, None, th)
        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)
        ok = [bool(t) and bool(same_reading(t, truth, disp, preset)) for t in texes]
        accepted = status == "auto_accepted"
        cases.append({"name": name, "truth": truth, "display": disp, "readings": reads, "ok": ok, "agree": agree,
                      "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)}}

    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>Ø 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())