#!/usr/bin/env python3
"""Self-contained ONNX inference for the HerBERT Polish legal NER model.
No PyTorch needed.  pip install onnxruntime transformers numpy

Run from the repo root:  python examples/inference_onnx.py
"""
import json
from pathlib import Path

import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer

ROOT = Path(__file__).resolve().parent.parent
PER_THRESHOLD = 0.2  # recall-first: flip a token to PER if summed PER prob >= this

tok = AutoTokenizer.from_pretrained(str(ROOT))
cfg = json.load(open(ROOT / "config.json", encoding="utf-8"))
id2label = {int(k): v for k, v in cfg["id2label"].items()}
label2id = cfg["label2id"]
sess = ort.InferenceSession(str(ROOT / "onnx" / "model_quantized.onnx"))
in_names = {i.name for i in sess.get_inputs()}


def softmax(x):
    e = np.exp(x - x.max(-1, keepdims=True))
    return e / e.sum(-1, keepdims=True)


def predict(text):
    enc = tok(text, return_offsets_mapping=True, return_tensors="np",
              truncation=True, max_length=512)
    offsets = enc["offset_mapping"][0]
    feeds = {"input_ids": enc["input_ids"].astype(np.int64),
             "attention_mask": enc["attention_mask"].astype(np.int64)}
    if "token_type_ids" in in_names:
        feeds["token_type_ids"] = np.zeros_like(enc["input_ids"], dtype=np.int64)
    probs = softmax(sess.run(None, feeds)[0][0])      # (seq, num_labels)
    ids = probs.argmax(-1)
    per_b, per_i = label2id["B-PER"], label2id["I-PER"]

    spans, cur = [], None  # cur = [start, end, type]
    for i, (s, e) in enumerate(offsets):
        if s == e:                                    # special token
            if cur:
                spans.append(cur); cur = None
            continue
        label = id2label[int(ids[i])]
        # recall-first override for persons
        if label == "O" and probs[i, per_b] + probs[i, per_i] >= PER_THRESHOLD:
            label = "I-PER" if (cur and cur[2] == "PER") else "B-PER"
        if label == "O":
            if cur:
                spans.append(cur); cur = None
            continue
        tag, etype = label.split("-", 1)
        if tag == "B" or cur is None or cur[2] != etype:
            if cur:
                spans.append(cur)
            cur = [int(s), int(e), etype]
        else:
            cur[1] = int(e)
    if cur:
        spans.append(cur)
    return [{"type": t, "start": a, "end": b, "text": text[a:b]} for a, b, t in spans]


if __name__ == "__main__":
    samples = [
        "Pozwany Jan Kowalski, zam. ul. Słoneczna 5 w Krakowie, PESEL 02070803628.",
        "Powódka Anna Nowak-Kowalska, e-mail a.nowak@example.pl, tel. 501 234 567.",
    ]
    for t in samples:
        print("\n" + t)
        for ent in predict(t):
            print(f"   {ent['type']:8} [{ent['start']:>3}:{ent['end']:<3}] {ent['text']!r}")