#!/usr/bin/env python3
"""Score the bench's task (a) with the estate's own cross-encoder reranker.

GAP 3 of `GAPS.md`, pre-registered before this file was written. It runs
**on this laptop's CPU**, from the copy of the model already in the local
Hugging Face cache — no GPU, no seat, and no box in the estate is called.

The output is `rows-gaps/rerank.a.jsonl`, one row per item carrying the item's
id, its label, the number of chunks the page was cut into, the max score and the
argmax chunk's heading. `tables_gaps.py` does the arithmetic; this file only
scores, so a re-run of the tables never needs the model again.

THE CASTING, verbatim from the pre-registration: the page's `state` is cut on
level-2 ATX headings, each section keeps its heading line, text before the first
heading is its own chunk, any section over 180 words is re-cut into 180-word
windows with 40 words of overlap, and the item's score is the MAX over chunks.
"""
from __future__ import annotations

import json
import os
import re
import time

HERE = os.path.dirname(os.path.abspath(__file__))
KIT = os.path.join(HERE, "kit")
OUT_DIR = os.path.join(HERE, "rows-gaps")
MODEL = "cross-encoder/ms-marco-MiniLM-L6-v2"
WINDOW, OVERLAP = 180, 40
HEADING = re.compile(r"^## +(.*)$", re.M)


def chunks(state: str) -> list[tuple[str, str]]:
    """(heading, text) for every chunk of one page, in document order."""
    cuts = [(m.start(), m.group(1).strip()) for m in HEADING.finditer(state)]
    spans: list[tuple[str, str]] = []
    if not cuts or cuts[0][0] > 0:
        head = state[: cuts[0][0]] if cuts else state
        if head.strip():
            spans.append(("(before the first heading)", head.strip()))
    for i, (pos, head) in enumerate(cuts):
        end = cuts[i + 1][0] if i + 1 < len(cuts) else len(state)
        spans.append((head, state[pos:end].strip()))
    out: list[tuple[str, str]] = []
    for head, text in spans:
        words = text.split()
        if len(words) <= WINDOW:
            out.append((head, text))
            continue
        step = WINDOW - OVERLAP
        for s in range(0, len(words), step):
            piece = words[s:s + WINDOW]
            if not piece:
                break
            out.append((head, " ".join(piece)))
            if s + WINDOW >= len(words):
                break
    return out


def main() -> int:
    import torch
    from sentence_transformers import CrossEncoder

    torch.set_num_threads(os.cpu_count() or 4)
    items = json.load(open(os.path.join(KIT, "task_a.json"), encoding="utf-8"))
    articles = json.load(open(os.path.join(KIT, "articles.json"), encoding="utf-8"))
    by_slug = {s: chunks(a["state"]) for s, a in articles.items()}
    print("pages %d, chunks %d, min/median/max per page %d/%d/%d"
          % (len(by_slug), sum(len(v) for v in by_slug.values()),
             min(len(v) for v in by_slug.values()),
             sorted(len(v) for v in by_slug.values())[len(by_slug) // 2],
             max(len(v) for v in by_slug.values())))

    t0 = time.monotonic()
    model = CrossEncoder(MODEL, device="cpu")
    load_s = time.monotonic() - t0
    os.makedirs(OUT_DIR, exist_ok=True)
    path = os.path.join(OUT_DIR, "rerank.a.jsonl")
    t1 = time.monotonic()
    with open(path, "w", encoding="utf-8") as fh:
        for it in items:
            ch = by_slug[it["slug"]]
            s0 = time.monotonic()
            scores = model.predict([(it["question"], c) for _, c in ch],
                                   batch_size=32, show_progress_bar=False)
            best = max(range(len(ch)), key=lambda i: float(scores[i]))
            fh.write(json.dumps({
                "id": it["id"], "slug": it["slug"], "label": it["label"],
                "kind": it["kind"], "question": it["question"],
                "n_chunks": len(ch), "score": float(scores[best]),
                "best_heading": ch[best][0],
                "seconds": time.monotonic() - s0,
            }) + "\n")
    total = time.monotonic() - t1
    print("model %s loaded in %.1f s; %d items scored in %.1f s on %d CPU threads"
          % (MODEL, load_s, len(items), total, torch.get_num_threads()))
    print("wrote", path)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())