formula_img2latex/tools/vergleich.py

#!/usr/bin/env python3
"""Tuning aid: verify a fixed sample of formulas and write comparison images + an HTML overview.

    python3 tools/vergleich.py AUSGABE [--n 50] [--seed 1] [--name gehrke-50]

The sample (formula ids) is stored in vergleich/<name>/sample.json and reused on later runs, so
changes to the procedure are compared on the same formulas. Output (not versioned, it shows
excerpts of the book):

    vergleich/<name>/
    ├── index.html         all formulas: image panel, LaTeX, measurements, decision, reason
    ├── summary.json       counts per status, all numbers
    ├── sample.json        the formula ids of the sample
    └── NN-<id>.png        Original / LaTeX rigid / LaTeX elastic / difference
"""

from __future__ import annotations

import argparse
import html
import io
import json
import random
import sys
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

from PIL import Image, ImageDraw, ImageFont

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

from formula_img2latex.checks import pandoc_mathml_ok  # noqa: E402
from formula_img2latex.decide import decide, thresholds  # noqa: E402
from formula_img2latex.pipeline import _distinct_candidates, _ref_for, _verify_job, open_all  # noqa: E402

FONT = "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"
LABELS = (("orig", "Original"), ("rigid", "LaTeX starr"), ("render", "LaTeX elastisch"), ("diff", "Differenz"))


def panel(pngs: dict[str, bytes], target_h: int = 70) -> Image.Image:
    """Stack the four images, enlarged with visible pixels, each with a label."""
    imgs = {k: Image.open(io.BytesIO(v)).convert("RGB") for k, v in pngs.items()}
    if not imgs:
        return Image.new("RGB", (300, 40), "white")
    h0 = next(iter(imgs.values())).height
    f = max(1, min(8, round(target_h / max(h0, 1))))
    try:
        font = ImageFont.truetype(FONT, 14)
    except OSError:
        font = ImageFont.load_default()
    rows = []
    for key, label in LABELS:
        if key not in imgs:
            continue
        im = imgs[key]
        im = im.resize((im.width * f, im.height * f), Image.NEAREST)
        row = Image.new("RGB", (im.width + 140, im.height + 8), "white")
        ImageDraw.Draw(row).text((4, 4), label, fill=(80, 80, 80), font=font)
        row.paste(im, (136, 4))
        rows.append(row)
    w = max(r.width for r in rows)
    out = Image.new("RGB", (w, sum(r.height for r in rows) + 4), "white")
    y = 0
    for r in rows:
        out.paste(r, (0, y))
        y += r.height
    ImageDraw.Draw(out).rectangle([0, 0, out.width - 1, out.height - 1], outline=(200, 200, 200))
    return out


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("source", type=Path)
    ap.add_argument("--n", type=int, default=50)
    ap.add_argument("--seed", type=int, default=1)
    ap.add_argument("--name", default=None, help="Unterordner in vergleich/ (Standard: <buch>-<n>)")
    ap.add_argument("--preset", default=None)
    ap.add_argument("--workers", type=int, default=16)
    ap.add_argument("--new-sample", action="store_true", help="neue Stichprobe ziehen")
    args = ap.parse_args()

    source, ws = open_all(args.source, None)
    name = args.name or f"{args.source.name.split('_')[0].lower()}-{args.n}"
    out = Path(__file__).resolve().parents[1] / "vergleich" / name
    out.mkdir(parents=True, exist_ok=True)
    sample_file = out / "sample.json"
    pool = [f for f in ws.all() if _distinct_candidates(f) and (f.classify or {}).get("label") != "figure"]
    if sample_file.exists() and not args.new_sample:
        ids = json.loads(sample_file.read_text())
    else:
        ids = sorted(f.id for f in random.Random(args.seed).sample(pool, min(args.n, len(pool))))
        sample_file.write_text(json.dumps(ids, indent=1))
    by_id = {f.id: f for f in pool}
    formulas = [by_id[i] for i in ids if i in by_id]
    preset = args.preset or ws.settings.get("preset") or "cm"
    book_ppt = ws.settings.get("book_px_per_pt")
    th = thresholds(ws.settings.get("thresholds"))

    jobs = [(f.id, source.original_bytes(_ref_for(source, f)), f.asset, _distinct_candidates(f), f.display, preset)
            for f in formulas]
    results: dict[str, tuple[dict, dict[str, bytes]]] = {}
    with ProcessPoolExecutor(max_workers=args.workers) as ex:
        for fut in as_completed([ex.submit(_verify_job, j) for j in jobs]):
            fid, res, pngs = fut.result()
            best = max(res, key=lambda r: (r["ok"], r["min_contrast"] if r["min_contrast"] is not None else 1.0,
                                           r["metrics"].get("ncc", 0.0)))
            results[fid] = (best, pngs)
            print(f"  {len(results)}/{len(jobs)}", file=sys.stderr, end="\r")
    print(file=sys.stderr)
    pandoc = dict(zip([f.id for f in formulas], pandoc_mathml_ok([(results[f.id][0]["tex"], f.display) for f in formulas])))

    for old in out.glob("*.png"):
        old.unlink()
    counts: dict[str, int] = {}
    rows_html, summary = [], []
    for n, f in enumerate(formulas, 1):
        best, pngs = results[f.id]
        status, reason = decide(best, pandoc[f.id], book_ppt, th)
        counts[status] = counts.get(status, 0) + 1
        img = f"{n:02d}-{f.id}.png"
        panel(pngs).save(out / img)
        m = best.get("metrics", {})
        mc = best.get("min_contrast")
        summary.append({"n": n, "id": f.id, "status": status, "reason": reason, "tex": best.get("tex"),
                        "ocr": [c.get("tex") for c in f.candidates if c.get("tex")], "display": f.display,
                        "min_contrast": mc, "worst_variant": best.get("worst_variant"),
                        "worst_kind": best.get("worst_kind"), "climbed_from": best.get("climbed_from"),
                        "metrics": m})
        climbed = ""
        if best.get("climbed_from"):
            climbed = "<div class=fix>korrigiert aus: <code>" + html.escape(best["climbed_from"][0]) + "</code></div>"
        rows_html.append(
            f"<section class='{status}'><h2>{n:02d} · <span class=st>{status}</span> · "
            f"{'abgesetzt' if f.display else 'im Text'} · {html.escape(f.occurrences[0]['file'])}:{f.occurrences[0]['line']}</h2>"
            f"<img src='{img}'><div><code>{html.escape(best.get('tex') or '')}</code></div>{climbed}"
            f"<div class=m>NCC elastisch {m.get('ncc', 0):.3f} · starr {m.get('rigid_ncc', 0):.3f} · "
            f"Verzerrung {m.get('warp_share', 0):.2f} · px/pt {m.get('px_per_pt', 0):.3f} · "
            f"Kontrast {('–' if mc is None else f'{mc:.3f}')} (knappste Variante: {html.escape(best.get('worst_kind') or '–')}) · "
            f"Pandoc {'ok' if pandoc[f.id] else 'FEHLER'}</div><div class=r>{html.escape(reason)}</div></section>")
    head = " · ".join(f"{k}: {v}" for k, v in sorted(counts.items()))
    (out / "index.html").write_text(
        "<!doctype html><meta charset=utf-8><title>Vergleich " + html.escape(name) + "</title><style>"
        "body{font:14px system-ui,sans-serif;margin:16px;background:#f6f6f3}section{background:#fff;border:1px solid #ddd;"
        "border-radius:6px;padding:8px 12px;margin:10px 0}h2{font-size:14px;margin:0 0 6px}img{max-width:100%}"
        ".auto_accepted .st{color:#1f7a3a}.needs_review .st{color:#a15c00}.rejected_keep_image .st,.render_failed .st{color:#a32020}"
        ".m{font:12px monospace;color:#555;margin-top:4px}.r{font-size:13px}.fix{font-size:12px;color:#2a5db0}"
        "code{font-size:12px}</style>"
        f"<h1>Vergleich {html.escape(name)}</h1><p>{len(formulas)} Formeln, Schrift {preset} · {head}</p>"
        "<p>Zeilen je Bild: Original · LaTeX starr ausgerichtet · LaTeX elastisch ausgerichtet · Differenz "
        "(schwarz = beide, rot = nur Original, blau = nur LaTeX, gelb = Abweichungsfleck)</p>"
        + "".join(rows_html), encoding="utf-8")
    (out / "summary.json").write_text(json.dumps({"preset": preset, "counts": counts, "formulas": summary},
                                                 ensure_ascii=False, indent=1), encoding="utf-8")
    print(f"{len(formulas)} Formeln geprüft – {head}")
    print(f"Übersicht: {out / 'index.html'}")
    return 0


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