core/generate/recognizer_anchor_inject.py

"""ADR-0163.D.2 — per-category recognizer anchor injection.

When the candidate-graph pipeline's existing parser yields no candidates
for a statement AND the ratified recognizer registry recognizes the
statement, this module is consulted to build typed solver primitives
(``CandidateInitial`` / future ``CandidateOperation`` values) from the
recognizer's ``parsed_anchors``.  The output extends ``per_sentence_choices``
the same way the existing parser's output does, so the downstream
solver runs unchanged.

Doctrine
--------
- Pure, deterministic injectors.  Same ``(match, sentence)`` → same
  ``SentenceChoice`` tuple, byte-equal.
- Refusal-preferring: each injector returns ``()`` when it cannot build
  a primitive that passes the existing ``_initial_admissible``
  structural check (the wrong=0 safety net the candidate-graph already
  enforces).
- No LLM / embeddings / learned classifiers; the injection is rules-only
  same discipline as Phase A/C/D detection.
- Per-category boundary: v1 implements only ``discrete_count_statement``.
  Every other category routes to the empty-tuple fallback (skip-only,
  identical to the round-2 Phase D wiring) and lands in follow-up
  D.2.x PRs after the framework's empirical lift is operator-reviewed.

Five-layer wrong=0 safety net (the Phase D.2 brief's load-bearing
section) is preserved across this module:

  1. Matcher narrowness — ``recognizer_match._try_extract_discrete_count_anchor``
     refuses on any ambiguity.
  2. Extraction correctness — anchor fields ground in the literal
     statement surface.
  3. Injection correctness — the per-category injector returns a
     ``CandidateInitial`` that passes ``_initial_admissible``; failure
     to ground yields ``()``.
  4. Replay gate — propose-time ``run_admissibility_replay_gate``
     auto-rejects any extraction change that lifts the GSM8K wrong
     count.
  5. Multi-branch decision rule — when an injected candidate disagrees
     with another branch's answer, the candidate-graph refuses.
"""

from __future__ import annotations

from typing import Mapping, Union

from evals.refusal_taxonomy.shape_categories import ShapeCategory
from generate.math_candidate_parser import CandidateInitial, CandidateOperation
from generate.math_problem_graph import (
    InitialPossession,
    MathGraphError,
    Operation,
    Quantity,
)
from generate.recognizer_match import RecognizerMatch

# ADR-0170 — the widened injector emission type. Per-category injectors
# may emit a tuple of ``CandidateInitial`` (existing) or
# ``CandidateOperation`` (new, ADR-0170). The downstream
# ``per_sentence_choices`` aggregator dispatches admissibility on the
# concrete type (``_initial_admissible`` vs ``roundtrip_admissible``).
# No new admission paths are introduced by the widening itself; new
# emission shapes ship in subsequent per-injector PRs (ADR-0170 §"impl
# outline" W2/W3/W4/W5).
InjectorEmission = Union[CandidateInitial, CandidateOperation]


# ---------------------------------------------------------------------------
# Public surface
# ---------------------------------------------------------------------------


def inject_from_match(
    match: RecognizerMatch,
    sentence: str,
) -> tuple[InjectorEmission, ...]:
    """Dispatch a recognizer match to its per-category injector.

    Returns an empty tuple when the category has no v1 injector or when
    the v1 injector refused. Per ADR-0170, the return type is now
    ``tuple[InjectorEmission, ...]`` (``CandidateInitial | CandidateOperation``)
    so per-category injectors can emit operations as well as initials.
    The v1 ``discrete_count_statement`` injector continues to emit only
    ``CandidateInitial`` — the widening is type-level only in this PR.
    """
    injector = _INJECTORS.get(match.category)
    if injector is None:
        return ()
    return injector(match, sentence)


# ---------------------------------------------------------------------------
# Per-category injectors
# ---------------------------------------------------------------------------


def inject_discrete_count_statement(
    match: RecognizerMatch,
    sentence: str,
) -> tuple[InjectorEmission, ...]:
    """Build CandidateInitial OR CandidateOperation from ``discrete_count``
    parsed anchors, dispatched on the matcher's ``anchor_kind``.

    Per ADR-0170 W2 — the matcher records ``anchor_kind`` as either
    ``"possession"`` (verbs ``has/have/had``) or ``"acquisition"``
    (verbs in ``_ACQUISITION_VERBS``).

    - ``possession`` → ``CandidateInitial`` (existing behavior; the
      sentence asserts an initial state)
    - ``acquisition`` → ``CandidateOperation(kind='add')`` (new in W2;
      the sentence asserts an add-operation, preserving
      ADR-0131.G.1's branch-disagreement discipline — the regex
      parser's ADD_VERBS path emits the same kind of operation for
      single-word units, so the injector path complements it on
      multi-word units without conflicting)

    v1 narrowness: at most one anchor per match; absent or
    unconstructable anchors return ``()``.
    """
    if not match.parsed_anchors:
        return ()

    out: list[InjectorEmission] = []
    for anchor in match.parsed_anchors:
        anchor_kind = anchor.get("anchor_kind", "possession")
        if anchor_kind == "possession":
            cand: InjectorEmission | None = _build_initial_from_discrete_count(
                anchor, sentence
            )
        elif anchor_kind == "acquisition":
            cand = _build_operation_from_discrete_count_acquisition(
                anchor, sentence
            )
        else:
            # Unknown anchor_kind — under-admit. Future widenings (e.g.
            # "depletion" verbs as CandidateOperation(subtract)) extend
            # this branch.
            return ()
        if cand is None:
            # Under-admit on any failure to construct.  Partial
            # admission would mean the downstream Cartesian product
            # enumerates a graph missing state.
            return ()
        out.append(cand)
    return tuple(out)


# ---------------------------------------------------------------------------
# Internals
# ---------------------------------------------------------------------------


def _build_initial_from_discrete_count(
    anchor: Mapping[str, object],
    sentence: str,
) -> CandidateInitial | None:
    """Construct one CandidateInitial from a discrete_count anchor.

    Refuses (returns ``None``) when any field cannot be coerced or when
    the constructed value would violate ``CandidateInitial`` /
    ``InitialPossession`` invariants.  The resulting CandidateInitial is
    structurally verified upstream by ``_initial_admissible``.

    Anchor schema:
        {
          "kind": "discrete_count",
          "subject_role": <str>,
          "count_token": <str>,         # '20' or 'two'
          "count_kind": <"integer"|"word">,
          "counted_noun": <str>,        # 'paperclips' / 'Pokemon cards'
        }
    """
    subject_role = anchor.get("subject_role")
    count_token = anchor.get("count_token")
    count_kind = anchor.get("count_kind")
    counted_noun = anchor.get("counted_noun")

    if (
        not isinstance(subject_role, str) or not subject_role
        or not isinstance(count_token, str) or not count_token
        or not isinstance(count_kind, str)
        or not isinstance(counted_noun, str) or not counted_noun
    ):
        return None

    # Resolve the count token to a numeric value.  v1 supports integer
    # and single-word cardinals; hyphenated compounds defer to a follow-up
    # PR because their resolution requires the language pack's
    # parse_compound_cardinal helper which is not on this hot path.
    value = _resolve_count_value(count_token, count_kind)
    if value is None:
        return None

    # CandidateInitial requires an anchor verb token recognized in its
    # post-init whitelist (has/have/had/owns/owned/holds/held/contains/
    # contained — matched by the recognizer's narrowness rule).  We pick
    # the literal verb token from the sentence so the round-trip ground
    # check inside _initial_admissible succeeds.  Falls back to 'has' when
    # the verb cannot be located in the surface; that fallback only fires
    # when the recognizer's match diverges from the sentence and is the
    # under-admit path.
    verb_in_sentence = _locate_possession_verb(sentence)
    if verb_in_sentence is None:
        return None

    try:
        quantity = Quantity(value=value, unit=counted_noun)
        initial = InitialPossession(entity=subject_role, quantity=quantity)
    except MathGraphError:
        return None

    try:
        return CandidateInitial(
            initial=initial,
            source_span=sentence,
            matched_anchor=verb_in_sentence,
            matched_value_token=count_token,
            matched_unit_token=counted_noun,
            matched_entity_token=subject_role,
        )
    except ValueError:
        return None


def _build_operation_from_discrete_count_acquisition(
    anchor: Mapping[str, object],
    sentence: str,
) -> CandidateOperation | None:
    """Construct one CandidateOperation(kind='add') from a discrete_count
    anchor whose ``anchor_kind == "acquisition"``.

    Per ADR-0170 W2 — acquisition verbs (``collected``, ``received``,
    ``bought``, ``got``) are routed to operations, not initials, in
    accordance with ADR-0131.G.1's branch-disagreement discipline. The
    solver's defaults-from-zero rule resolves single-statement
    acquisitions correctly (``0 + N = N``).

    Refuses (returns ``None``) when any field cannot be coerced, when
    the literal verb token cannot be located in the surface, or when
    the constructed ``CandidateOperation`` would violate its post-init
    invariants. The result is admissibility-checked upstream by
    ``roundtrip_admissible``.

    Anchor schema (same as possession, with ``anchor_kind`` discriminator):
        {
          "kind": "discrete_count",
          "anchor_kind": "acquisition",
          "subject_role": <str>,
          "count_token": <str>,
          "count_kind": <"integer"|"word">,
          "counted_noun": <str>,
          "verb_token": <str>,  # e.g. "collected"
        }
    """
    subject_role = anchor.get("subject_role")
    count_token = anchor.get("count_token")
    count_kind = anchor.get("count_kind")
    counted_noun = anchor.get("counted_noun")
    verb_token = anchor.get("verb_token")

    if (
        not isinstance(subject_role, str) or not subject_role
        or not isinstance(count_token, str) or not count_token
        or not isinstance(count_kind, str)
        or not isinstance(counted_noun, str) or not counted_noun
        or not isinstance(verb_token, str) or not verb_token
    ):
        return None

    value = _resolve_count_value(count_token, count_kind)
    if value is None:
        return None

    # Locate the literal verb surface in the sentence so the
    # round-trip ground check in ``roundtrip_admissible`` succeeds.
    # The matcher already confirmed ``verb_token`` is in
    # ``_ACQUISITION_VERBS`` (which is itself a subset of
    # ``ADD_VERBS``), so the downstream CandidateOperation post-init
    # whitelist accepts the matched_verb token.
    located_verb = _locate_token(sentence, verb_token)
    if located_verb is None:
        return None

    try:
        operand = Quantity(value=value, unit=counted_noun)
        op = Operation(
            actor=subject_role,
            kind="add",
            operand=operand,
        )
    except MathGraphError:
        return None

    try:
        return CandidateOperation(
            op=op,
            source_span=sentence,
            matched_verb=located_verb,
            matched_value_token=count_token,
            matched_unit_token=counted_noun,
            matched_actor_token=subject_role,
        )
    except ValueError:
        return None


def _locate_token(sentence: str, target_lc: str) -> str | None:
    """Return the literal-surface form of ``target_lc`` (lowercased) in
    ``sentence`` whitespace-tokenized, or ``None`` if absent.

    Used by the acquisition-verb path to extract the matched verb
    surface for ``CandidateOperation.matched_verb``. Falls back to
    ``None`` only when the matcher's recorded ``verb_token`` somehow
    diverges from the sentence surface — the under-admit path.
    """
    for raw in sentence.split():
        tok = raw.strip(".,;:!?\"'()[]{}").lower()
        if tok == target_lc:
            return tok
    return None


def _resolve_count_value(count_token: str, count_kind: str) -> int | None:
    """Map ``count_token`` to a numeric value.

    Integer tokens parse with ``int``.  Word-form tokens look up
    ``WORD_NUMBERS`` from the language pack; unknown words refuse.
    Hyphenated compounds (``twenty-five``) defer to D.2.x — v1 returns
    ``None`` for them.
    """
    if count_kind == "integer":
        try:
            return int(count_token)
        except ValueError:
            return None
    if count_kind == "word":
        # Local import to keep module import-time cheap and to avoid a
        # circular import via the math_candidate_parser surface.
        from generate.math_roundtrip import WORD_NUMBERS

        token_lc = count_token.lower()
        if token_lc in WORD_NUMBERS:
            return int(WORD_NUMBERS[token_lc])
        # Hyphenated compound: defer to D.2.x.
        return None
    return None


def _locate_possession_verb(sentence: str) -> str | None:
    """Return the first possession-anchor verb (lowercased) found in
    ``sentence`` whitespace-tokenized, or ``None`` when absent.

    The verb is the surface token that ``CandidateInitial.__post_init__``
    validates against its registered anchor whitelist.  Returning the
    LITERAL surface keeps the round-trip ground check in
    ``_initial_admissible`` honest.
    """
    possession_verbs = ("has", "have", "had")
    for raw in sentence.split():
        tok = raw.strip(".,;:!?\"'()[]{}").lower()
        if tok in possession_verbs:
            return tok
    return None


# ---------------------------------------------------------------------------
# Dispatch table — keep deterministic and explicit.
# Adding a category here is the SINGLE place a new D.2.x category
# registers its injector.  No global state, no side effects.
# ---------------------------------------------------------------------------

_INJECTORS: Mapping[ShapeCategory, "type"] = {
    ShapeCategory.DISCRETE_COUNT_STATEMENT: inject_discrete_count_statement,  # type: ignore[dict-item]
    # All other recognizer categories route to the empty-tuple fallback
    # in ``inject_from_match`` — `_INJECTORS.get(category)` returns
    # ``None`` and the dispatcher returns ``()``, which the
    # candidate-graph then treats as "recognizer matched but produced
    # no injection" → explicit refusal (the wrong=0 fix from #359).
    #
    # Categories deferred to follow-up PRs:
    #
    # ShapeCategory.DESCRIPTIVE_SETUP_NO_QUANTITY — by design (no quantity)
    # ShapeCategory.RATE_WITH_CURRENCY           — needs CandidateRate
    #                                              (SentenceChoice union
    #                                              extension; ADR-0171)
    # ShapeCategory.TEMPORAL_AGGREGATION         — needs apply_rate primitive
    #                                              in the algebra
    # ShapeCategory.MULTIPLICATIVE_AGGREGATION   — emits
    #                                              CandidateInitial(product)
    #                                              after ADR-0170 widens
    #                                              return type
    # ShapeCategory.CURRENCY_AMOUNT              — A1 currency_amount;
    #                                              CandidateInitial-shaped,
    #                                              ships after ADR-0170
    #
    # See docs/decisions/ADR-0170-injector-contract-widening.md for the
    # contract widening that unblocks DCS-S1 / A1 / A3.
}


__all__ = [
    "InjectorEmission",
    "inject_from_match",
    "inject_discrete_count_statement",
]