"""QUE-259 full per-region diagnostic.

For every table region docling finds, answer two things from independent sources:

  (1) Did OCR run for this region?  -> from docling's own parsed-page TextCells
      (from_ocr + confidence) that fall inside the table bbox. No quber geometry
      beyond point-in-box; the OCR verdict is docling's, not inferred.

  (2) Is docling's content comparable to Camelot+LLM's?  -> put docling's
      OCR-reconstructed table (rows x cols, markdown) next to the matched
      Camelot+LLM ExtractedTable (accuracy, markdown) for the same region.

Run:  uv run python experiments/que259/diagnostic.py <doc_id> [<doc_id> ...]
Needs: tables/<id>.tables.json already present (the Camelot+LLM / SoM output).
Re-parses the PDF with generate_parsed_pages=True to recover from_ocr.
"""

from __future__ import annotations

import json
import sys
from pathlib import Path

from docling.datamodel.accelerator_options import AcceleratorDevice
from quber.core.parsers.docling_parser import parser_for_preset

HERE = Path(__file__).parent
OUT = HERE / "out"
OUT.mkdir(exist_ok=True)


def tl_from_table_bbox(bb, H):
    """docling table bbox (BOTTOMLEFT) -> TOPLEFT (x0,y0,x1,y1)."""
    return (bb.l, H - bb.t, bb.r, H - bb.b)


def tl_from_cell_rect(rect):
    xs = [rect.r_x0, rect.r_x1, rect.r_x2, rect.r_x3]
    ys = [rect.r_y0, rect.r_y1, rect.r_y2, rect.r_y3]
    return (min(xs), min(ys), max(xs), max(ys))


def center_in(box, region) -> bool:
    cx, cy = (box[0] + box[2]) / 2, (box[1] + box[3]) / 2
    return region[0] <= cx <= region[2] and region[1] <= cy <= region[3]


def overlap_cov(a, b):
    ix1, iy1 = max(a[0], b[0]), max(a[1], b[1])
    ix2, iy2 = min(a[2], b[2]), min(a[3], b[3])
    inter = max(0.0, ix2 - ix1) * max(0.0, iy2 - iy1)
    aa = max(1e-9, (a[2] - a[0]) * (a[3] - a[1]))
    return inter / aa


def nonempty_cells(table):
    return sum(1 for c in table.data.table_cells if (c.text or "").strip())


def run(doc_id, parser):
    po = parser.build_pipeline_options()
    po.generate_parsed_pages = True
    res = parser.build_converter(po).convert(str(HERE / "inputs" / f"{doc_id}.pdf"))
    doc = res.document
    # page-indexed OCR cells (TOPLEFT boxes)
    page_cells = {}
    page_size = {}
    for i, pg in enumerate(res.pages, start=1):
        if pg.size:
            page_size[i] = (pg.size.width, pg.size.height)
        page_cells[i] = [(tl_from_cell_rect(c.rect), c.from_ocr, c.confidence or 0.0) for c in pg.cells]

    som = json.loads((HERE / "tables" / f"{doc_id}.tables.json").read_text())

    rows = []
    for t in doc.tables:
        if not t.prov:
            continue
        p = t.prov[0].page_no
        W, H = page_size.get(p, (612.0, 792.0))
        region = tl_from_table_bbox(t.prov[0].bbox, H)
        cells = [c for c in page_cells.get(p, []) if center_in(c[0], region)]
        n = len(cells)
        ocr_frac = (sum(c[1] for c in cells) / n) if n else 0.0
        conf = (sum(c[2] for c in cells if c[1]) / max(1, sum(c[1] for c in cells))) if n else 0.0
        # match a Camelot+LLM table on same page by bbox overlap (use som_region in pts)
        best, best_cov = None, 0.0
        for st in som:
            if st["page"] != p:
                continue
            sr = st.get("som_region")
            sbox = (sr[0] * W, sr[1] * H, sr[2] * W, sr[3] * H) if sr else None
            cov = overlap_cov(region, sbox) if sbox else 0.0
            if cov > best_cov:
                best, best_cov = st, cov
        cam_md = (best.get("markdown") if best else "") or ""
        rows.append({
            "doc": doc_id, "page": p,
            "ocr_ran": ocr_frac > 0.5, "ocr_cells": n, "ocr_frac": round(ocr_frac, 2), "ocr_conf": round(conf, 2),
            "docling_rows": t.data.num_rows, "docling_cols": t.data.num_cols,
            "docling_nonempty_cells": nonempty_cells(t),
            "docling_md": t.export_to_markdown(doc),
            "camelot_match_cov": round(best_cov, 2),
            "camelot_accuracy": (best.get("camelot_accuracy") if best else None),
            "camelot_nonempty": bool(cam_md.strip()),
            "camelot_md": cam_md,
        })
    return rows


def verdict(r):
    if r["ocr_ran"] and not r["camelot_nonempty"]:
        return "OCR-only (docling recovered, Camelot empty)"
    if not r["ocr_ran"] and r["camelot_nonempty"]:
        return "text-layer (both can read)"
    if r["ocr_ran"] and r["camelot_nonempty"]:
        return "both produced output"
    return "neither produced output"


def main():
    ids = sys.argv[1:] or ["101230425", "102174103"]
    parser = parser_for_preset("tuned-financial", accelerator_device=AcceleratorDevice.AUTO)
    all_rows = []
    for doc_id in ids:
        all_rows.extend(run(doc_id, parser))
    json.dump(all_rows, open(OUT / "diagnostic.json", "w"), indent=1)

    lines = ["# QUE-259 per-region diagnostic: OCR determination + Camelot+LLM comparability", ""]
    lines.append(f"{'doc':11} {'pg':>3} {'ocr_ran':>7} {'frac':>5} {'conf':>5} {'dRxC':>7} {'dCells':>6} {'cam_acc':>7} {'cam?':>5}  verdict")
    for r in all_rows:
        lines.append(
            f"{r['doc']:11} {r['page']:>3} {str(r['ocr_ran']):>7} {r['ocr_frac']:>5} {r['ocr_conf']:>5} "
            f"{str(r['docling_rows'])+'x'+str(r['docling_cols']):>7} {r['docling_nonempty_cells']:>6} "
            f"{str(r['camelot_accuracy']):>7} {str(r['camelot_nonempty']):>5}  {verdict(r)}"
        )
    lines.append("")
    # one paired example per verdict type
    seen = set()
    for r in all_rows:
        v = verdict(r)
        if v in seen:
            continue
        seen.add(v)
        lines.append(f"## Example -- {v}  ({r['doc']} p{r['page']})")
        lines.append(f"OCR ran: {r['ocr_ran']} (frac={r['ocr_frac']}, conf={r['ocr_conf']}, {r['ocr_cells']} cells)")
        lines.append(f"docling reconstruction: {r['docling_rows']}x{r['docling_cols']}, {r['docling_nonempty_cells']} non-empty cells")
        lines.append("```\n" + (r["docling_md"][:600] or "(empty)") + "\n```")
        lines.append(f"Camelot+LLM: accuracy={r['camelot_accuracy']} nonempty={r['camelot_nonempty']} (match cov={r['camelot_match_cov']})")
        lines.append("```\n" + (r["camelot_md"][:600] or "(empty)") + "\n```")
        lines.append("")
    report = "\n".join(lines)
    (OUT / "diagnostic.md").write_text(report)
    print(report)


if __name__ == "__main__":
    main()