#!/usr/bin/env python3
import regex as re
import torch


TOKEN_RE = re.compile(r"[A-Za-z0-9]+|[^\w\s]", re.UNICODE)


def tokenize_with_spans(text: str):
    return [(m.group(0), m.start(), m.end()) for m in TOKEN_RE.finditer(text)]


def ppsn_label_ids(model) -> list[int]:
    ids = []
    for raw_id, raw_label in model.config.id2label.items():
        label_id = int(raw_id)
        label = str(raw_label or "").strip()
        if label.endswith("PPSN"):
            ids.append(label_id)
    return sorted(ids)


def word_aligned_ppsn_spans(text: str, model, tokenizer, threshold: float) -> list[dict]:
    pieces = tokenize_with_spans(text)
    if not pieces:
        return []

    words = [word for word, _, _ in pieces]
    encoded = tokenizer(words, is_split_into_words=True, return_tensors="pt", truncation=True)
    word_ids = encoded.word_ids(batch_index=0)

    device = next(model.parameters()).device
    encoded = {k: v.to(device) for k, v in encoded.items()}

    with torch.no_grad():
        logits = model(**encoded).logits[0]

    probs = torch.softmax(logits, dim=-1)
    label_ids = ppsn_label_ids(model)

    word_scores: list[float] = []
    for word_index in range(len(pieces)):
        score = 0.0
        for token_index, wid in enumerate(word_ids):
            if wid != word_index:
                continue
            for label_id in label_ids:
                score = max(score, float(probs[token_index, label_id]))
        word_scores.append(score)

    spans: list[dict] = []
    active = None
    for (word, start, end), score in zip(pieces, word_scores):
        if score >= threshold:
            if active is None:
                active = {"start": start, "end": end, "score": score}
            else:
                active["end"] = end
                active["score"] = max(active["score"], score)
        elif active is not None:
            spans.append(active)
            active = None

    if active is not None:
        spans.append(active)

    for span in spans:
        span["text"] = text[span["start"] : span["end"]]
        span["label"] = "PPSN"
        span["source"] = "model"

    return spans