#!/usr/bin/env python3
"""
furniture.py — detect recurring page chrome by cross-page repetition; derive content margins;
reclassify pictures (real figure vs barcode/QR/header furniture). The first-pass mask.

Principle: an item at ~the same (x,y) on many pages is **chrome, not question content**. This
needs no classifier — pure positional recurrence — and it solves the genuine gap the overlay
surfaced (the app-generated QR top-right and the foot barcode being mislabelled context_figure),
including the QR that bleeds past the margin. It also yields the content margins so stage-2 analysis
can be fed only the question/response region.

Outputs a mask + margins JSON, and an A/B summary (figure false-positives before vs after masking).

Usage:
  python furniture.py <docling_doc.json> [--freq 0.4] [--out results/furniture.json]
"""
import json, argparse
from collections import defaultdict

GRID = 24          # pt — position quantisation; items sharing a cell across pages are "recurring"


def gather(doc):
    out = []
    for key in ("texts", "pictures", "tables"):
        for it in doc.get(key, []):
            prov = it.get("prov") or []
            bb = prov[0].get("bbox") if prov else None
            pg = prov[0].get("page_no") if prov else None
            if bb and pg:
                out.append({"page": pg, "kind": key[:-1], "label": it.get("label", key[:-1]),
                            "bbox": bb, "text": (it.get("text") or "")[:40]})
    return out


def cell(bb):
    return (round((bb["l"] + bb["r"]) / 2 / GRID), round((bb["t"] + bb["b"]) / 2 / GRID))


def detect(items, n_pages, freq):
    """Flag each item furniture=True if its position-cell appears on >= freq*n_pages pages."""
    pages_at = defaultdict(set)
    for it in items:
        pages_at[cell(it["bbox"])].add(it["page"])
    fcells = {c: len(p) for c, p in pages_at.items() if len(p) >= freq * n_pages}
    for it in items:
        it["furniture"] = cell(it["bbox"]) in fcells
    return fcells


def content_margins(items):
    """Content x-band + per-page content bbox from NON-furniture items (what stage-2 should see)."""
    body = [it for it in items if not it["furniture"]]
    if not body:
        return None
    lefts = sorted(it["bbox"]["l"] for it in body)
    rights = sorted(it["bbox"]["r"] for it in body)
    band = {"x_left": round(lefts[max(0, len(lefts) // 20)], 1),       # 5th pct — robust to strays
            "x_right": round(rights[min(len(rights) - 1, len(rights) * 19 // 20)], 1)}
    per_page = {}
    bp = defaultdict(list)
    for it in body:
        bp[it["page"]].append(it["bbox"])
    for pg, bbs in bp.items():
        per_page[pg] = {"top": round(max(b["t"] for b in bbs), 1),
                        "bottom": round(min(b["b"] for b in bbs), 1),
                        "left": round(min(b["l"] for b in bbs), 1),
                        "right": round(max(b["r"] for b in bbs), 1)}
    return {"content_x_band": band, "per_page": per_page}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("doc")
    ap.add_argument("--freq", type=float, default=0.40, help="recurrence fraction => furniture")
    ap.add_argument("--out", default="results/furniture.json")
    a = ap.parse_args()
    doc = json.load(open(a.doc))
    items = gather(doc)
    n_pages = len({it["page"] for it in items})
    fcells = detect(items, n_pages, a.freq)
    margins = content_margins(items)

    pics = [it for it in items if it["kind"] == "picture"]
    pics_furn = [it for it in pics if it["furniture"]]
    txt_furn = [it for it in items if it["kind"] == "text" and it["furniture"]]
    # break furniture pictures down by cell (which recurring object)
    by_cell = defaultdict(list)
    for it in pics_furn:
        by_cell[cell(it["bbox"])].append(it)

    result = {
        "n_pages": n_pages, "freq_threshold": a.freq,
        "furniture_cells": {f"{c[0]},{c[1]}": n for c, n in sorted(fcells.items())},
        "content_margins": margins,
        "ab_test_figures": {
            "context_figure_before_mask": len(pics),
            "context_figure_after_mask": len(pics) - len(pics_furn),
            "removed_as_furniture": len(pics_furn),
            "removed_breakdown": {f"cell {c[0]},{c[1]}": len(v) for c, v in sorted(by_cell.items())},
        },
        "text_furniture_removed": len(txt_furn),
        "items": items,   # each carries furniture flag — consumed by overlay.py --furniture
    }
    json.dump(result, open(a.out, "w"))

    ab = result["ab_test_figures"]
    print(f"pages {n_pages}  freq>={a.freq}  furniture cells: {result['furniture_cells']}")
    print(f"content x-band: {margins['content_x_band'] if margins else None}")
    print(f"\nA/B — figure (picture) classification:")
    print(f"  context_figure BEFORE mask : {ab['context_figure_before_mask']}")
    print(f"  context_figure AFTER  mask : {ab['context_figure_after_mask']}")
    print(f"  removed as furniture       : {ab['removed_as_furniture']}  {ab['removed_breakdown']}")
    print(f"  text furniture removed     : {result['text_furniture_removed']} (page numbers / 'Turn over' / headers)")
    print(f"-> wrote {a.out}")


if __name__ == "__main__":
    main()
