#!/usr/bin/env python3
"""Build the gliner2 adjudication queue: Luna-labeled spans rescored with
the singlepass bundle (rafmacalaba/gliner-datause-catchall-singlepass).

Source: the gliner2 config of rafmacalaba/datause-ner (passage rows with
spans[]: text, luna_label, char start/end, key). Local mirrors in
hf_datause_ner/gliner2_*.jsonl are byte-identical to the Hub (verified:
29,346/29,346 keys) and used as the fetch source.

Per mention the queue carries:
    luna            camp2 verdict, 1=keep / 0=drop
    head_score      probe_score from the singlepass infer head (MPS)
    extractor_score GLiNER proposer score @0.1 for the same grid cell
                    (null when the proposer didn't fire on the span)
    band            keep/confusion/drop via probe_labels.decide

Sampling: round-robin over origins, multi-mention passages first, span
budget --limit (default 520). Overlapping duplicate spans are kept as
separate mentions (the UI renderer suppresses in-text doubles).

    uv run python human_labeling/build_gliner2_queue.py [--limit 520] [--batch 8]
"""

import argparse
import json
import sys
from collections import defaultdict
from pathlib import Path

HERE = Path(__file__).resolve().parent
REPO = HERE.parent
sys.path.insert(0, str(REPO))

MIRROR = REPO / "hf_datause_ner"
OUT = HERE / "queue_gliner2.json"


def load_passages() -> list[dict]:
    rows = []
    for split in ("train", "val", "holdout"):
        for line in (MIRROR / f"gliner2_{split}.jsonl").read_text().splitlines():
            if line.strip():
                r = json.loads(line)
                r["split"] = split
                rows.append(r)
    return rows


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="rafmacalaba/gliner-datause-catchall-singlepass")
    ap.add_argument("--limit", type=int, default=520, help="span budget")
    ap.add_argument("--batch", type=int, default=8)
    a = ap.parse_args()

    import torch
    from training.singlepass_infer import default_device, load_bundle
    from training.probe_features_infer import char_to_infer_word
    from probe_labels import decide

    passages = load_passages()
    # round-robin origins, multi-mention passages first within each origin
    by_origin: dict[str, list[dict]] = defaultdict(list)
    for p in passages:
        by_origin[p["origin"]].append(p)
    ordered: list[dict] = []
    origins = sorted(by_origin)
    for o in origins:
        by_origin[o].sort(key=lambda p: -min(len(p["spans"]), 2))
    i = 0
    while any(by_origin[o] for o in origins):
        o = origins[i % len(origins)]
        if by_origin[o]:
            ordered.append(by_origin[o].pop(0))
        i += 1

    # take passages until span budget; dedupe exact-duplicate keys in passage
    chosen, n_spans = [], 0
    for p in ordered:
        if n_spans >= a.limit:
            break
        seen, spans = set(), []
        for s in p["spans"]:
            if s["key"] in seen:
                continue
            seen.add(s["key"])
            spans.append(s)
        chosen.append((p, spans))
        n_spans += len(spans)

    device = default_device()
    print(f"device={device} model={a.model} passages={len(chosen)} spans={n_spans}",
          flush=True)
    model, head, bundle = load_bundle(
        "rafmacalaba/gliner-datause-mentions-catch-all", a.model, device)
    thresholds = bundle.get("thresholds") or {}
    radius = bundle["radius"]

    import torch.utils.data
    INFER_LABELS = ["DATA_MENTION"]
    texts = [p["input"] for p, _ in chosen]
    prepared = model.prepare_batch(texts, INFER_LABELS)
    collator = model.create_collator()

    def collate_fn(batch):
        return model.collate_batch(batch, prepared["entity_types"], collator)

    loader = torch.utils.data.DataLoader(
        prepared["input_x"], batch_size=a.batch, shuffle=False,
        collate_fn=collate_fn)

    v2o = prepared["valid_to_orig_idx"]
    o2v = {o: v for v, o in enumerate(v2o)}
    n_probe = n_ext = n_skip = 0
    items: list[dict] = []

    def flush(p, spans, probes, extractors):
        mentions = []
        for s, probe, ext in zip(spans, probes, extractors):
            mentions.append({
                "key": s["key"], "surface": s["text"],
                "start": s["start"], "end": s["end"],
                "luna": s.get("luna_label"),
                "head_score": probe, "extractor_score": ext,
                "band": decide(probe, p["origin"], thresholds),
            })
        bands = {m["band"] for m in mentions}
        pband = ("unscored" if "unscored" in bands else
                 "confusion" if "confusion" in bands else
                 "mixed" if len(bands) > 1 else bands.pop())
        items.append({
            "queue": "gliner2", "origin": p["origin"], "split": p["split"],
            "ctx": p["input"], "n": len(mentions), "band": pband,
            "mentions": mentions, "scored_by": a.model,
        })

    row = 0
    with torch.no_grad():
        for batch in loader:
            out = model.run_batch(batch, threshold=0.1, move_to_device=True)
            W = out.words_embedding.detach().float()
            mask = (out.mask.detach().cpu()
                    if getattr(out, "mask", None) is not None else None)
            decoded = model.decode_batch(out, batch, threshold=0.1,
                                         flat_ner=True, multi_label=False)
            B = W.shape[0]
            for bi in range(B):
                vi = row + bi
                oi = v2o[vi]
                p, spans = chosen[oi]
                if o2v.get(oi) is None:  # empty passage filtered upstream
                    flush(p, spans, [None] * len(spans), [None] * len(spans))
                    n_skip += len(spans)
                    continue
                w = W[bi].to(device)
                L = int(mask[bi].sum()) if mask is not None else w.shape[0]
                starts = prepared["start_token_map"][vi]
                proposals = [(int(sp.start), int(sp.end), float(sp.score))
                             for sp in decoded[bi]]
                probes, extractors = [], []
                for s in spans:
                    cs, ce = s["start"], s["end"]
                    g0, g1 = char_to_infer_word(starts, cs, ce)
                    probe = ext = None
                    if g1 < L and g0 < L:
                        idx = torch.arange(g0, g1 + 1, device=device)
                        parts = [w[g0], w[g1], w[idx].mean(dim=0)]
                        if radius > 0:
                            w0, w1 = max(0, g0 - radius), min(g1 + radius, L - 1)
                            parts.append(w[w0:w1 + 1].mean(dim=0))
                        probe = float(torch.sigmoid(
                            head(torch.cat(parts).unsqueeze(0))).item())
                        n_probe += 1
                        hit = [sc for (ps, pe, sc) in proposals
                               if ps == g0 and pe == g1]
                        if hit:
                            ext = hit[0]
                            n_ext += 1
                    probes.append(probe)
                    extractors.append(ext)
                flush(p, spans, probes, extractors)
            row += B

    OUT.write_text("\n".join(json.dumps(it) for it in items) + "\n")
    from collections import Counter
    bands = Counter(m["band"] for it in items for m in it["mentions"])
    lu = Counter(m["luna"] for it in items for m in it["mentions"])
    print(f"queue: passages={len(items)} spans={sum(it['n'] for it in items)} "
          f"probe_scored={n_probe} extractor_matched={n_ext} "
          f"unscored={n_skip} bands={dict(bands)} luna={dict(lu)} -> {OUT}",
          flush=True)


if __name__ == "__main__":
    main()