"""Span pointer decoding for the v3 copy head."""
from __future__ import annotations

import torch


def best_span(start: torch.Tensor, end: torch.Tensor, max_len: int = 32):
    """Highest p(start) * p(end) with start <= end < start + max_len -> (s, e, prob)."""
    ps, pe = start.float().softmax(-1), end.float().softmax(-1)
    Q = ps.shape[0]
    i = torch.arange(Q, device=ps.device)
    ok = (i[None, :] >= i[:, None]) & (i[None, :] < i[:, None] + max_len)
    score = (ps[:, None] * pe[None, :]).masked_fill(~ok, 0)
    flat = int(score.argmax())
    return flat // Q, flat % Q, float(score.flatten()[flat])


def _cls(ch: str) -> str:
    return "a" if ch.isalpha() else "d" if ch.isdigit() else ""


def snap_word(text: str, s: int, e: int) -> tuple[int, int]:
    """Widen the char span [s, e) so it never cuts a word or number in half ("US Doll" -> "US Dollars").
    Only letter-letter and digit-digit edges are extended, so "100" stays "100" inside "100km"."""
    while s > 0 and s < e and _cls(text[s]) and _cls(text[s - 1]) == _cls(text[s]):
        s -= 1
    while e < len(text) and e > s and _cls(text[e - 1]) and _cls(text[e]) == _cls(text[e - 1]):
        e += 1
    return s, e