"""PropBank-style semantic-role label vocabulary for SRL supervision.

The relation space is *shared* with UD: UD deprels keep ids ``[0, NUM_DEPREL)``
(unchanged), and SRL roles occupy the contiguous block starting at ``SRL_BASE ==
NUM_DEPREL``. SRL losses operate on the ``edge_logits`` slice
``[SRL_BASE : SRL_BASE + NUM_SRL_ROLES]`` using **local** role ids ``[0,
NUM_SRL_ROLES)`` where local id 0 is ``NONE`` (this ordered pair is not an
argument edge). A model that does SRL therefore needs
``graph_relation_types >= RELATION_VOCAB_SIZE``.
"""

from __future__ import annotations

from strata.data.ud_labels import IGNORE_INDEX, NUM_DEPREL  # noqa: F401 (re-export IGNORE_INDEX)

# local id 0 == NONE (no argument edge); remaining are PropBank roles.
SRL_ROLES: tuple[str, ...] = (
    "NONE",
    "ARG0", "ARG1", "ARG2", "ARG3", "ARG4", "ARG5", "ARGA",
    "ARGM-TMP", "ARGM-LOC", "ARGM-MNR", "ARGM-CAU", "ARGM-DIS", "ARGM-ADV",
    "ARGM-MOD", "ARGM-NEG", "ARGM-PRP", "ARGM-DIR", "ARGM-EXT", "ARGM-PRD",
    "ARGM-GOL", "ARGM-COM", "ARGM-REC", "ARGM-LVB", "ARGM-OTHER", "ARG-OTHER",
)
SRL_ROLE_TO_LOCAL: dict[str, int] = {role: i for i, role in enumerate(SRL_ROLES)}
NONE_LOCAL = 0
NUM_SRL_ROLES = len(SRL_ROLES)

SRL_BASE = NUM_DEPREL  # offset into the shared relation vocabulary
RELATION_VOCAB_SIZE = SRL_BASE + NUM_SRL_ROLES

# Column markers in Universal Propositions role columns.
PREDICATE_MARKER = "V"
NO_ROLE = "_"


def srl_role_to_local(role: str) -> int:
    """Map a UP role string to a local SRL id (never ``NONE`` for a real role).

    Continuation/reference prefixes (``C-`` / ``R-``) are dropped to their base;
    unrecognised core / modifier roles fall back to generic buckets.
    """

    r = role.strip().upper()
    if r.startswith(("C-", "R-")):
        r = r[2:]
    if r in SRL_ROLE_TO_LOCAL and r != "NONE":
        return SRL_ROLE_TO_LOCAL[r]
    if r.startswith("ARGM"):
        return SRL_ROLE_TO_LOCAL["ARGM-OTHER"]
    if r.startswith("ARG"):
        return SRL_ROLE_TO_LOCAL["ARG-OTHER"]
    return SRL_ROLE_TO_LOCAL["ARG-OTHER"]