#!/usr/bin/env python3
"""Time and accuracy per recognition tool on formulas whose TeX is known (serial, one at a time).

    python3 tools/zeittest.py EPUB [--n 15] [--tools local,local2,cloud,p2t-cpu,p2t-gpu] [--gpu cuda:0]

Same sample as `calibrate --consensus` (seed 42, real images). Every tool reads every formula alone,
one request after the other, after one uncounted warm-up on another formula. Writes
vergleich/zeittest/index.html (not versioned): table per tool, then per formula the original, the
true formula and every reading set with pdflatex, with time and verdict; plus ergebnis.json.
"""

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.consensus import same_reading  # noqa: E402
from formula_img2latex.engines import make_engine  # noqa: E402
from formula_img2latex.engines.p2t import Pix2TextEngine  # noqa: E402
from formula_img2latex.render import render_many  # noqa: E402
from formula_img2latex.tex import strip_delimiters  # noqa: E402

LABELS = {"local": "Qwen3-VL-30B (lokal, GPU)", "local2": "Gemma-4-26B (lokal, GPU)", "cloud": "GLM-4.6V (Cloud)",
          "p2t-cpu": "pix2text MFR 1.5 (CPU)", "p2t-gpu": "pix2text MFR 1.5 (GPU)"}


def engine_for(tool: str, gpu: str):
    if tool == "p2t-cpu":
        return Pix2TextEngine(device="cpu")
    if tool == "p2t-gpu":
        return Pix2TextEngine(device=gpu)
    return make_engine(tool)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("epub", type=Path)
    ap.add_argument("--n", type=int, default=15)
    ap.add_argument("--tools", default="local,local2,cloud,p2t-cpu,p2t-gpu")
    ap.add_argument("--gpu", default="cuda", help="Gerät für p2t-gpu (Index relativ zu CUDA_VISIBLE_DEVICES)")
    ap.add_argument("--name", default="zeittest")
    args = ap.parse_args()
    tools = [t.strip() for t in args.tools.split(",") if t.strip()]

    items = corpus(args.epub)
    smp = random.Random(42).sample(items, min(150, len(items)))   # as calibrate_consensus
    formulas, warm = smp[:args.n], smp[args.n]
    preset = _best_preset(smp[:12])

    runs: dict[str, dict] = {}
    for tool in tools:
        eng = engine_for(tool, args.gpu)
        t0 = time.time()
        eng.check()
        load = time.time() - t0
        eng.recognize(warm[1], warm[0], "", warm[3])   # warm-up, not counted
        reads = []
        for name, data, truth, disp in formulas:
            r = eng.recognize(data, name, "", disp)
            ok = bool(r.tex) and bool(same_reading(r.tex, truth, disp, preset))
            reads.append({"tex": r.tex, "seconds": r.seconds, "ok": ok, "error": r.error,
                          "tokens_in": r.tokens_in, "tokens_out": r.tokens_out})
            print(f"{tool:8s} {r.seconds:6.2f}s {'✓' if ok else '✗'} {r.tex or r.error}", file=sys.stderr, flush=True)
        secs = [x["seconds"] for x in reads]
        runs[tool] = {"label": LABELS.get(tool, tool), "model": eng.model, "load_s": load, "reads": reads,
                      "mean_s": statistics.mean(secs), "median_s": statistics.median(secs),
                      "correct": sum(x["ok"] for x in reads), "errors": sum(bool(x["error"]) for x in reads),
                      "tokens_in": sum(x["tokens_in"] for x in reads), "tokens_out": sum(x["tokens_out"] for x in reads)}

    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, "truth": [f[2] for f in formulas], "runs": runs},
        ensure_ascii=False, indent=1), encoding="utf-8")

    n = len(formulas)
    rows = "".join(
        f"<tr><td>{html.escape(r['label'])}</td><td>{r['mean_s']:.2f}</td><td>{r['median_s']:.2f}</td>"
        f"<td>{r['correct']}/{n} ({r['correct'] / n:.0%})</td><td>{r['load_s']:.1f}</td>"
        f"<td>{r['tokens_in'] + r['tokens_out'] or '–'}</td></tr>" for r in runs.values())
    parts = [f"<h1>Zeittest: {html.escape(args.epub.name[:70])}</h1>",
             f"<p>{n} Formeln mit bekanntem TeX (alt-Text), jedes Tool seriell, eine Formel nach der anderen, "
             "nach einem nicht gezählten Aufwärmlauf. Alle Formeln mit pdflatex gesetzt. ✓ = gleiche Zeichen an "
             "gleicher Stelle wie die wahre Formel (Abstände, kursiv/aufrecht, Fettdruck egal).</p>",
             "<table class=sum><tr><th>Tool</th><th>Ø s/Formel</th><th>Median s</th><th>richtig</th>"
             f"<th>Laden s</th><th>Tokens</th></tr>{rows}</table>"]

    texes = []
    for i, (_n, _d, truth, disp) in enumerate(formulas):
        texes.append((strip_delimiters(truth), disp))
        texes += [(runs[t]["reads"][i]["tex"] or "\\text{–}", disp) for t in tools]
    rendered = render_many(texes, preset)
    per = 1 + len(tools)
    for i, (name, data, truth, _disp) in enumerate(formulas):
        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>")
        body = [f"<tr><th>Original</th><td><img src='{png_uri(orig_img(data, 2))}'></td><td></td></tr>",
                f"<tr><th>richtig</th><td>{imgs[0]}</td><td><code>{html.escape(truth)}</code></td></tr>"]
        for k, t in enumerate(tools):
            r = runs[t]["reads"][i]
            mark = "<b class=ok>✓</b>" if r["ok"] else "<b class=no>✗</b>"
            body.append(f"<tr><th>{mark} {html.escape(runs[t]['label'])}<br><small>{r['seconds']:.2f} s</small></th>"
                        f"<td>{imgs[k + 1]}</td><td><code>{html.escape(r['tex'] or r['error'] or '–')}</code></td></tr>")
        parts.append(f"<section><div class=v>Formel {i + 1} · {html.escape(name)}</div><table>{''.join(body)}</table></section>")
    (out / "index.html").write_text(
        "<!doctype html><meta charset=utf-8><title>Zeittest</title><style>"
        "body{font:14px system-ui,sans-serif;margin:16px;background:#f6f6f3}section{background:#fff;border:1px solid #ddd;"
        "border-radius:6px;padding:6px 10px;margin:10px 0}.v{font-weight:600;font-size:13px;margin-bottom:4px}"
        "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:210px;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(f"\n| Tool | Ø s/Formel | Median s | richtig (n={n}) | Laden s |")
    print("|---|---|---|---|---|")
    for r in runs.values():
        print(f"| {r['label']} | {r['mean_s']:.2f} | {r['median_s']:.2f} | {r['correct']}/{n} | {r['load_s']:.1f} |")
    if "p2t-cpu" in runs and "p2t-gpu" in runs:
        same = sum(a["tex"] == b["tex"] for a, b in zip(runs["p2t-cpu"]["reads"], runs["p2t-gpu"]["reads"]))
        print(f"\npix2text CPU und GPU gleicher Text: {same}/{n}")
    print(out / "index.html")
    return 0


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