"""Score image-to-Word outputs against the benchmark's ground truth.

Layout: results/<tool>/<sample>.docx  (sample = 1-report, 2-table, ...)
Run:    python3 -I docs/outreach/bench/score.py [--json | --ts]
        --ts regenerates apps/web/src/lib/benchImageToWord.ts (the page's data)
        and copies this script to apps/web/public/bench/score.py.txt

Everything is read from the .docx XML, nothing is judged by eye:
  text      1 - character error rate against the page's text (math excluded on
            both sides: equations are scored by count, and their linear text
            differs legitimately between converters). Order-sensitive: a page
            whose columns come back interleaved scores low even if every word
            is right — which is what an editor of that file experiences.
  words     share of the page's words that appear in the output, order ignored
            (multiset recall). Separates "misread" from "out of order".
  tables    rows x cols of each w:tbl (grid columns, after merges)
  merged    cells with gridSpan > 1 or a vMerge restart
  equations native Word equations (m:oMath)
  headings  paragraphs styled Title / Heading n
  lists     paragraphs carrying list numbering
  image     True when the document is a picture with (almost) no text in it
"""

import json
import re
import shutil
import sys
import unicodedata
import zipfile
from pathlib import Path

HERE = Path(__file__).resolve().parent
W = "{http://schemas.openxmlformats.org/wordprocessingml/2006/main}"

import xml.etree.ElementTree as ET


def norm(s: str) -> str:
    s = unicodedata.normalize("NFKC", s)
    s = re.sub(r"[‐-―−]", "-", s)
    s = re.sub(r"[‘’]", "'", s).replace("“", '"').replace("”", '"')
    return re.sub(r"\s+", " ", s).strip()


def cer(ref: str, hyp: str) -> float:
    if not ref:
        return 0.0
    prev = list(range(len(hyp) + 1))
    for i, a in enumerate(ref, 1):
        cur = [i]
        for j, b in enumerate(hyp, 1):
            cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (a != b)))
        prev = cur
    return prev[-1] / len(ref)


def words_recall(ref: str, hyp: str) -> float:
    from collections import Counter
    tok = lambda t: re.findall(r"[a-z0-9]+", t.lower())
    r, h = Counter(tok(ref)), Counter(tok(hyp))
    total = sum(r.values())
    return sum(min(c, h[w]) for w, c in r.items()) / total if total else 0.0


def para_text(p) -> str:
    # w:t only — m:t (math) is deliberately left out, see module docstring.
    return "".join(t.text or "" for t in p.iter(W + "t"))


def read_docx(path: Path) -> dict:
    with zipfile.ZipFile(path) as z:
        names = z.namelist()
        doc = ET.fromstring(z.read("word/document.xml"))
        part = lambda kind: [ET.fromstring(z.read(n)) for n in sorted(names) if re.match(rf"word/{kind}\d*\.xml$", n)]
        headers, footers = part("header"), part("footer")
        raw = z.read("word/document.xml").decode("utf8", "replace")
        # Style ids are localised (a German Word writes "Titel", "berschrift1");
        # the w:name of a built-in style is always the English one, so match on it.
        names = {}
        if "word/styles.xml" in z.namelist():
            for st in ET.fromstring(z.read("word/styles.xml")).iter(W + "style"):
                n = st.find(W + "name")
                names[st.get(W + "styleId")] = (n.get(W + "val") if n is not None else "") or ""

    def block(root) -> str:
        return " ".join(para_text(p) for p in root.iter(W + "p"))

    # Reading order of a page: header, body, footer. A file may carry several
    # header/footer variants (first page, even, default) holding the same line;
    # each distinct one counts once, or packaging would cost points twice.
    seen = []
    for t in (block(r) for r in [*headers]):
        if norm(t) and norm(t) not in seen:
            seen.append(norm(t))
    tail = []
    for t in (block(r) for r in footers):
        if norm(t) and norm(t) not in tail:
            tail.append(norm(t))
    text = norm(" ".join([*seen, block(doc), *tail]))

    tables = []
    merged = 0
    for tbl in doc.iter(W + "tbl"):
        rows = tbl.findall(W + "tr")
        cols = len(tbl.findall(f"{W}tblGrid/{W}gridCol")) or max((len(r.findall(W + "tc")) for r in rows), default=0)
        tables.append([len(rows), cols])
        # Direct cells only: iter() would also walk a nested table's cells and
        # count them again for every table that encloses it.
        for tc in (tc for tr in rows for tc in tr.findall(W + "tc")):
            pr = tc.find(W + "tcPr")
            if pr is None:
                continue
            span = pr.find(W + "gridSpan")
            vm = pr.find(W + "vMerge")
            if span is not None and int(span.get(W + "val", "1")) > 1:
                merged += 1
            elif vm is not None and vm.get(W + "val") == "restart":
                merged += 1

    headings = 0
    lists = 0
    for p in doc.iter(W + "p"):
        pr = p.find(W + "pPr")
        if pr is None:
            continue
        st = pr.find(W + "pStyle")
        if st is not None and re.match(r"(?i)(title|heading\s*\d)$", names.get(st.get(W + "val", ""), st.get(W + "val", ""))):
            headings += 1
        if pr.find(W + "numPr") is not None:
            lists += 1

    return {
        "text": text,
        "tables": tables,
        "merged": merged,
        "equations": raw.count("<m:oMath>") + raw.count("<m:oMath "),
        "headings": headings,
        "lists": lists,
        "pictures": raw.count("<w:drawing") + raw.count("<w:pict"),
    }


ROOT = HERE.parent.parent.parent
TS = ROOT / "apps/web/src/lib/benchImageToWord.ts"
KEYS = [("tool", "tool"), ("sample", "sample"), ("textAccuracy", "text_accuracy"), ("wordRecall", "word_recall"),
        ("tables", "tables"), ("tablesExpected", "tables_expected"), ("merged", "merged"),
        ("mergedExpected", "merged_expected"), ("equations", "equations"),
        ("equationsExpected", "equations_expected"), ("headings", "headings"),
        ("headingsExpected", "headings_expected")]


def write_ts(rows: list) -> None:
    """The page's data module. Rows for samples a tool has no .docx for are left
    out (the page renders those cells as "—"), never written as empty rows."""
    head = [
        "// GENERATED by `python3 -I docs/outreach/bench/score.py --ts` — do not edit by hand.",
        "// Rerun the benchmark and regenerate, so the page cannot drift from the data.",
        "",
        "export interface BenchRow {",
        "  tool: string;",
        "  sample: string;",
        "  textAccuracy: number;",
        "  wordRecall: number;",
        "  tables: number[][];",
        "  tablesExpected: number[][];",
        "  merged: number;",
        "  mergedExpected: number;",
        "  equations: number;",
        "  equationsExpected: number;",
        "  headings: number;",
        "  headingsExpected: number;",
        "}",
        "",
        "export const BENCH_ROWS: BenchRow[] = [",
    ]
    body = ["  " + json.dumps({k: r[src] for k, src in KEYS}) + "," for r in rows]
    TS.write_text("\n".join(head + body + ["];", ""]))
    shutil.copyfile(Path(__file__), ROOT / "apps/web/public/bench/score.py.txt")
    print(f"wrote {TS.relative_to(ROOT)} ({len(rows)} rows) and public/bench/score.py.txt")


def main() -> None:
    truth = json.loads((HERE / "samples" / "truth.json").read_text())
    rows = []
    for tool_dir in sorted(p for p in (HERE / "results").iterdir() if p.is_dir() and any(p.glob("*.docx"))):
        for sample, t in truth.items():
            f = tool_dir / f"{sample}.docx"
            if not f.exists():
                rows.append({"tool": tool_dir.name, "sample": sample, "missing": True})
                continue
            d = read_docx(f)
            ref = norm(t["text"])
            acc = max(0.0, 1 - cer(ref, d["text"]))
            rows.append(
                {
                    "tool": tool_dir.name,
                    "sample": sample,
                    "text_accuracy": round(acc, 3),
                    "word_recall": round(words_recall(ref, d["text"]), 3),
                    "tables": d["tables"],
                    "tables_expected": t["expect"]["tables"],
                    "merged": d["merged"],
                    "merged_expected": t["expect"]["merged"],
                    "equations": d["equations"],
                    "equations_expected": t["expect"]["equations"],
                    "headings": d["headings"],
                    "headings_expected": t["expect"]["headings"],
                    "lists": d["lists"],
                    "image_only": d["pictures"] > 0 and len(d["text"]) < 0.2 * len(ref),
                }
            )
    if "--ts" in sys.argv:
        write_ts([r for r in rows if not r.get("missing")])
        return
    if "--json" in sys.argv:
        print(json.dumps(rows, indent=2, ensure_ascii=False))
        return
    hdr = f"{'tool':<16}{'sample':<12}{'text':>6}{'words':>7}  {'tables (want)':<22}{'merged':>8}{'eq':>7}{'head':>6}{'list':>6}  note"
    print(hdr)
    for r in rows:
        if r.get("missing"):
            print(f"{r['tool']:<16}{r['sample']:<12}{'—':>6}  (no output)")
            continue
        tb = f"{r['tables']} ({r['tables_expected']})"
        print(
            f"{r['tool']:<16}{r['sample']:<12}{r['text_accuracy']:>6.1%}{r['word_recall']:>7.1%}  {tb:<22}"
            f"{r['merged']:>4}/{r['merged_expected']:<3}{r['equations']:>3}/{r['equations_expected']:<3}"
            f"{r['headings']:>6}{r['lists']:>6}  {'IMAGE ONLY' if r['image_only'] else ''}"
        )


if __name__ == "__main__":
    main()