#!/usr/bin/env python3
"""
bands.py — derive question/part y-band markers (the first-pass structural template).

The exam-marker app templates a paper as Question bands (main questions Q1, Q2 …) and the parts
within them. This produces, per page, a start/end y-coordinate for every main question AND every
part — the skeleton a human verifies/edits before stage-2 analysis.

Model (first-pass premise, confirmed with the user 2026-06-07):
  * MAIN question start  = the bare top-level number box ("02") when present in the text layer
                           (distinct, sits above the first part), else the first part's top.
  * PART start           = the part label's top (we already carry this geometry).
  * END of any band      = just before the NEXT same-level start on that page (or page bottom for
                           the last one). Parts are nested: a part's end never exceeds its question's.
Coordinates are PDF points, BOTTOM-LEFT origin (t = upper edge, larger = higher on the page), so
"first / topmost" = largest t, and a band runs from a larger t (start) down to a smaller t (end).

Usage:
  python bands.py <structured.json> [--docling results/E_tess_full.json] [--out results/bands/x.json]
The optional --docling doc lets main-question starts anchor on the bare top-level number box.
"""
import json, re, glob, argparse
from collections import defaultdict

LABEL_COL_MAX = 80           # left x-band where the boxed question/part numbers live


def _topnumber_boxes(docs):
    """{(page, qint): t} — bare top-level number boxes ('02') in the left label column, scanned
    across one or more Docling docs. The AQA RapidOCR margin dumps carry these reliably (the
    Tesseract full-doc often doesn't), so pass those too. Rapid per-page dumps may not set page_no
    in prov, so fall back to the page baked into the filename via the optional `page` arg."""
    out = {}
    for doc, page_hint in docs:
        for it in doc.get("texts", []):
            prov = it.get("prov") or []
            bb = prov[0].get("bbox") if prov else None
            pg = (prov[0].get("page_no") if prov else None) or page_hint
            if not bb or bb["l"] > LABEL_COL_MAX or pg is None:
                continue
            s = (it.get("text") or "").strip().replace(" ", "")
            m = re.match(r"^(\d{1,2})$", s)
            if m:
                key = (pg, int(m.group(1)))
                out[key] = max(bb["t"], out.get(key, bb["t"]))   # header box sits high (largest t)
    return out


def _ends(items):
    """Given [(key, start_t, extra...)] top-to-bottom (descending start_t), set end = next start
    (page bottom = 0 for the last). Returns list of dicts with start/end."""
    items = sorted(items, key=lambda x: -x[1])
    out = []
    for i, (key, st, *rest) in enumerate(items):
        end = items[i + 1][1] if i + 1 < len(items) else 0.0
        out.append((key, st, end, rest))
    return out


def derive_bands(result, doc=None, rapid_glob=None):
    docs = []
    if doc:
        docs.append((doc, None))
    for fn in sorted(glob.glob(rapid_glob) if rapid_glob else []):
        m = re.search(r"p(\d+)\.json", fn)
        docs.append((json.load(open(fn)), int(m.group(1)) if m else None))
    topnum = _topnumber_boxes(docs)
    # gather parts with geometry, grouped by page
    by_page = defaultdict(list)               # page -> [(q, label, t, b)]
    for q in result.get("questions", []):
        for p in q["parts"]:
            bb, pg = p.get("bbox"), p.get("page")
            if bb and pg:
                by_page[pg].append((q["question"], p["label"], bb["t"], bb["b"]))

    # global first page each question appears on (to mark the true start vs continuation pages)
    q_first_page = {}
    for pg, parts in by_page.items():
        for q, *_ in parts:
            q_first_page[q] = min(pg, q_first_page.get(q, pg))

    pages = {}
    for pg, parts in by_page.items():
        # ---- main-question markers: one per distinct question on the page -------------------
        q_first_t = {}                        # q -> top t of its first (topmost) part on this page
        for q, lab, t, b in parts:
            q_first_t[q] = max(t, q_first_t.get(q, t))
        main_starts = []
        for q, ft in q_first_t.items():
            tn = topnum.get((pg, int(re.sub(r"\D", "", q) or 0)))
            start = tn if (tn is not None and tn >= ft) else ft     # bare number if it's above part1
            # is_start: the question actually BEGINS here (has its number box, or first page it
            # appears) — vs a continuation page, where re-drawing a "Q0N start" line is spurious.
            is_start = (tn is not None) or (pg == q_first_page.get(q))
            main_starts.append((q, start, is_start))
        main = [{"question": q, "y_start": round(st, 1), "y_end": round(en, 1),
                 "is_start": rest[0]}
                for (q, st, en, rest) in _ends(main_starts)]
        main_band = {m["question"]: (m["y_start"], m["y_end"]) for m in main}

        # ---- part markers: each part label top; end = next part start, clipped to its question -
        part_items = [((q, lab), t) for q, lab, t, b in parts]
        part = []
        for (q, lab), st, en, _ in _ends(part_items):
            qen = main_band.get(q, (st, 0))[1]                      # don't run past the question end
            part.append({"label": lab, "question": q,
                         "y_start": round(st, 1), "y_end": round(max(en, qen), 1)})
        pages[pg] = {"main": main, "part": part}

    return {"board": result.get("board"), "paper_code": result.get("paper_code"),
            "coord_origin": "BOTTOMLEFT", "pages": pages}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("structured")
    ap.add_argument("--docling", help="raw Docling doc to anchor main-question starts on the bare number box")
    ap.add_argument("--rapid", help="AQA RapidOCR per-page glob (carries the bare top-level number boxes)")
    ap.add_argument("--out", default="results/bands.json")
    a = ap.parse_args()
    res = json.load(open(a.structured))
    doc = json.load(open(a.docling)) if a.docling else None
    bands = derive_bands(res, doc, a.rapid)
    json.dump(bands, open(a.out, "w"), indent=2)
    nq = sum(len(p["main"]) for p in bands["pages"].values())
    npt = sum(len(p["part"]) for p in bands["pages"].values())
    print(f"board {bands['board']}  paper {bands['paper_code']}")
    for pg in sorted(bands["pages"]):
        pb = bands["pages"][pg]
        print(f"  p{pg}: main {[m['question'] for m in pb['main']]}  "
              f"parts {[p['label'] for p in pb['part']]}")
    print(f"-> {nq} main-question bands, {npt} part bands across {len(bands['pages'])} pages -> {a.out}")


if __name__ == "__main__":
    main()
