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

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

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. The formulas are processed exactly
like `run` does (consensus mode: Qwen, Gemma, GLM for agreed readings, image veto). Output is
not versioned because it shows excerpts of the book:

    vergleich/<name>/index.html   per formula: original, rendered readings, decision, veto overlay
    vergleich/<name>/summary.json counts and all readings
"""

from __future__ import annotations

import argparse
import base64
import html
import io
import json
import random
import sys
from pathlib import Path

import numpy as np
from PIL import Image

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

from formula_img2latex.pipeline import _wanted, open_all, original_bytes, readings_of, run_consensus  # noqa: E402
from formula_img2latex.render import render_many  # noqa: E402
from formula_img2latex.tex import strip_delimiters  # noqa: E402

NAMES = {"local": "Qwen", "local2": "Gemma", "cloud": "GLM"}
COLOR = {"auto_accepted": "#1f7a3a", "needs_review": "#a15c00"}


def uri(img: Image.Image) -> str:
    buf = io.BytesIO()
    img.save(buf, "PNG")
    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()


def original(data: bytes, target_h: int = 60) -> Image.Image:
    im = Image.open(io.BytesIO(data))
    if im.mode in ("P", "LA", "RGBA") or "transparency" in im.info:
        im = im.convert("RGBA")
        im = Image.alpha_composite(Image.new("RGBA", im.size, "white"), im)
    im = im.convert("RGB")
    f = max(1, min(6, round(target_h / max(im.height, 1))))
    return im.resize((im.width * f, im.height * f), Image.NEAREST)


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)
    ap.add_argument("--new-sample", action="store_true")
    ap.add_argument("--no-cloud", action="store_true")
    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 _wanted(f, False)]
    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))

    run_consensus(source, ws, chunk=len(ids), ids=ids, use_cloud=not args.no_cloud)

    formulas = [ws.load(i) for i in ids]
    formulas = [f for f in formulas if f is not None]
    preset = ws.settings.get("preset") or "cm"
    items, rows = [], []
    for f in formulas:
        for r in readings_of(f):
            items.append((strip_delimiters(r.get("tex") or "") or "\\text{–}", f.display))
    rendered = iter(render_many(items, preset))
    counts: dict[str, int] = {}
    summary = []
    for n, f in enumerate(formulas, 1):
        counts[f.status] = counts.get(f.status, 0) + 1
        orig = original(original_bytes(source, ws, f))
        lines = [f"<tr><th>Original</th><td><img src='{uri(orig)}'></td><td></td></tr>"]
        reads = []
        for r in readings_of(f):
            rr = next(rendered)
            img = ""
            if rr.ok:
                z = 2.0 * max(1.0, orig.height / 60)
                g = (rr.rasterize(min(z, 4.0)) * 255).astype(np.uint8)
                img = uri(Image.fromarray(255 - g))
            who = NAMES.get(r.get("engine", ""), r.get("engine", "?"))
            reads.append({"engine": who, "model": r.get("model"), "tex": r.get("tex"), "p_mean": r.get("p_mean")})
            p = r.get("p_mean")
            lines.append(f"<tr><th>{who}</th><td>{f'<img src={chr(39)}{img}{chr(39)}>' if img else '(nicht renderbar)'}</td>"
                         f"<td><code>{html.escape(r.get('tex') or '–')}</code>"
                         f"{f' <small>p={p:.2f}</small>' if p is not None else ''}</td></tr>")
        diff = ws.review_dir / f"{f.id}-diff.png"
        if diff.exists():
            d = Image.open(diff).convert("RGB")
            d = d.resize((d.width * 2, d.height * 2), Image.NEAREST)
            lines.append(f"<tr><th>Veto-Overlay</th><td><img src='{uri(d)}'></td><td><small>rot = nur Original, "
                         "blau = nur LaTeX, schwarz = beide</small></td></tr>")
        color = COLOR.get(f.status, "#a32020")
        occ = f.occurrences[0] if f.occurrences else {}
        rows.append(f"<section style='border-left-color:{color}'><h2>{n:02d} · <span style='color:{color}'>{f.status}</span> · "
                    f"{'abgesetzt' if f.display else 'im Text'} · {html.escape(occ.get('file', ''))}:{occ.get('line', '')}</h2>"
                    f"<div class=r>{html.escape(f.reason or '')}</div><table>{''.join(lines)}</table></section>")
        summary.append({"n": n, "id": f.id, "status": f.status, "reason": f.reason, "final_tex": f.final_tex,
                        "readings": reads})
    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-left:6px solid #999;border-radius:6px;padding:8px 12px;margin:10px 0}h2{font-size:14px;margin:0 0 4px}"
        "th{text-align:left;font-weight:400;color:#666;width:100px;vertical-align:middle}td{padding:2px 8px;vertical-align:middle}"
        ".r{font-size:13px;margin-bottom:4px}code{font-size:12px}img{max-width:760px}</style>"
        f"<h1>Vergleich {html.escape(name)}</h1><p>{len(formulas)} Formeln · Regel: Qwen = Gemma = GLM, Fettdruck egal, "
        f"Bild-Veto · {head}</p>" + "".join(rows), encoding="utf-8")
    (out / "summary.json").write_text(json.dumps({"counts": counts, "formulas": summary}, ensure_ascii=False, indent=1),
                                      encoding="utf-8")
    print(f"{len(formulas)} Formeln – {head}")
    print(f"Übersicht: {out / 'index.html'}")
    return 0


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