"""QUE-259 -- discriminate tabular images (table regions with no text layer).

Method under test:
  A region that docling identifies as a TABLE (its TableFormer fires, via OCR on
  the page image) but that has ~0 native text in the source PDF underneath it is
  a *tabular image* -- a table that exists only as pixels. Camelot, which reads
  the PDF text layer, can never extract it.

Signals, all from artifacts we already produce -- no new model, no OCR of our own:
  S1  native_text_chars : chars in the source PDF under the region bbox (PyMuPDF)
  S2  docling_label     : TABLE vs PICTURE(class) that docling assigned the region
  S3  page_has_text     : whether the whole page has any native text layer

Classification per docling region:
  text_table     : docling TABLE, native text present            (camelot can read)
  tabular_image  : docling TABLE, native text ~0                  (image of a table)
  chart_image    : docling PICTURE classed chart/plot, text ~0    (not a table)
  other_image    : docling PICTURE (logo/photo/...), text ~0
"""

from __future__ import annotations

import glob
import json
import os

import fitz  # PyMuPDF

HERE = os.path.dirname(__file__)
INPUTS = os.path.join(HERE, "inputs")
DOCLING = os.path.join(HERE, "docling")
OUT = os.path.join(HERE, "out")
os.makedirs(OUT, exist_ok=True)

CHART = ("chart", "plot", "histogram", "heatmap", "graph", "diagram")
TEXT_EMPTY = 5  # chars under a region at/below this == no native text layer


def top_class(anns):
    for a in anns or []:
        if "predicted_classes" in a:
            cl = a["predicted_classes"]
            if cl:
                t = max(cl, key=lambda c: c.get("confidence", 0))
                return t.get("class_name"), round(t.get("confidence", 0), 3)
    return None


def rect_from_bottomleft(bbox, page_h):
    """docling bbox {l,t,r,b BOTTOMLEFT} -> fitz Rect (TOPLEFT, y-down)."""
    l, t, r, b = bbox["l"], bbox["t"], bbox["r"], bbox["b"]
    return fitz.Rect(l, page_h - t, r, page_h - b)


def chars_in(page, rect) -> int:
    return len(page.get_text("text", clip=rect).strip())


def analyze(doc_id):
    d = json.load(open(os.path.join(DOCLING, f"{doc_id}.docling.json")))
    pdf = fitz.open(os.path.join(INPUTS, f"{doc_id}.pdf"))
    page_text = {i + 1: len(pdf[i].get_text("text").strip()) for i in range(len(pdf))}
    regions = []

    def handle(item, kind):
        if not item.get("prov"):
            return
        prov = item["prov"][0]
        pg = prov["page_no"]
        if pg < 1 or pg > len(pdf):
            return
        page = pdf[pg - 1]
        rect = rect_from_bottomleft(prov["bbox"], page.rect.height)
        chars = chars_in(page, rect)
        cls = top_class(item.get("annotations")) if kind == "picture" else None
        is_chart = bool(cls and any(tok in cls[0] for tok in CHART))
        if kind == "table":
            label = "text_table" if chars > TEXT_EMPTY else "tabular_image"
        elif is_chart:
            label = "chart_image"
        else:
            label = "other_image"
        regions.append({
            "doc": doc_id, "page": pg, "kind": kind, "docling_class": cls,
            "native_text_chars": chars, "page_text_chars": page_text[pg],
            "bbox_bl": [round(v, 1) for v in (prov["bbox"]["l"], prov["bbox"]["t"], prov["bbox"]["r"], prov["bbox"]["b"])],
            "label": label,
        })

    for t in d.get("tables", []):
        handle(t, "table")
    for p in d.get("pictures", []):
        handle(p, "picture")
    pdf.close()
    return regions


def main():
    ids = sorted(os.path.basename(f).replace(".docling.json", "") for f in glob.glob(os.path.join(DOCLING, "*.docling.json")))
    all_regions = []
    for doc_id in ids:
        all_regions.extend(analyze(doc_id))
    json.dump(all_regions, open(os.path.join(OUT, "regions.json"), "w"), indent=1)

    # aggregate
    from collections import Counter, defaultdict
    by_label = Counter(r["label"] for r in all_regions)
    per_doc = defaultdict(Counter)
    for r in all_regions:
        per_doc[r["doc"]][r["label"]] += 1

    lines = ["# QUE-259 tabular-image discriminator -- results", ""]
    lines.append(f"Regions analyzed: {len(all_regions)} across {len(ids)} docs")
    lines.append(f"Label totals: {dict(by_label)}")
    lines.append("")
    lines.append(f"{'doc':12} {'text_table':>10} {'tabular_image':>13} {'chart_image':>11} {'other_image':>11}")
    for doc_id in ids:
        c = per_doc[doc_id]
        lines.append(f"{doc_id:12} {c['text_table']:>10} {c['tabular_image']:>13} {c['chart_image']:>11} {c['other_image']:>11}")
    lines.append("")

    # the separation evidence: native text under docling TABLE regions
    tab = [r for r in all_regions if r["kind"] == "table"]
    img_tab = [r["native_text_chars"] for r in tab if r["label"] == "tabular_image"]
    txt_tab = [r["native_text_chars"] for r in tab if r["label"] == "text_table"]
    def stats(xs):
        if not xs:
            return "none"
        xs = sorted(xs)
        return f"n={len(xs)} min={xs[0]} median={xs[len(xs)//2]} max={xs[-1]}"
    lines.append("Native text chars under docling TABLE regions (the discriminator):")
    lines.append(f"  tabular_image (flagged): {stats(img_tab)}")
    lines.append(f"  text_table   (passed):   {stats(txt_tab)}")
    lines.append("")

    # mixed docs: prove per-region (not per-doc) behavior
    lines.append("Mixed docs (both text and image tables present):")
    for doc_id in ids:
        c = per_doc[doc_id]
        if c["text_table"] and c["tabular_image"]:
            lines.append(f"  {doc_id}: {c['text_table']} text_table + {c['tabular_image']} tabular_image")
    report = "\n".join(lines)
    open(os.path.join(OUT, "discriminator.md"), "w").write(report)
    print(report)


if __name__ == "__main__":
    main()