stack/src/pfs/reaction.py

"""Public-reaction series per code family (``pfs.code_reaction``, #690,
task 7 of the 2026-09-09 code-family-anchors plan).

Two period kinds, one ``ReactionRow`` shape:

* ``"docket"`` — one row per regulations.gov docket (2017-2026, ten
  sealed dockets, ``bib.Store.dockets()``): how many distinct
  commenters (``item_key``) mentioned the family, out of the docket's
  total commenters, straight from the ``comments`` pgvector
  collection's chunk metadata (``docket``, ``item_key``, ``codes``,
  ``families``) — never from ``Docket.counts_json``, which counts raw
  comment documents, not distinct commenters who named the family.
  When ``--stance-sample N`` is passed, the newest *N* commenters' first
  matching chunk is classified into ``STANCES`` by the self-hosted
  closed-vocabulary classifier (``llm.classify``); an answer outside
  the vocabulary counts as ``"unclear"``.
* ``"fr-pairs"`` — one row per FR rule item, back to 2001: how many
  ``Comment:``/``Response:`` paragraph pairs in that rule's
  ``fr_anchors`` name a family code (``n_items``), out of the rule's
  total ``Comment:`` paragraph count — family or not (``n_total``) —
  the pre-2017 proxy, since regulations.gov comment text itself isn't
  indexed for rules that old.

``series`` builds both halves for one family (or one ``--code``
pseudo-family, a single-code family keyed by the code itself) and
returns them as one list, dockets first (docket id order) then FR pairs
(rule item_key order); the CLI (``stack pfs reaction``) prints and
optionally writes them.
"""

from __future__ import annotations

import json
import re
from collections import Counter
from typing import Any, Callable, Sequence

from sqlalchemy import text

from pfs.anchors import code_or_family_predicate
from pfs.descriptors import rule_year_of
from pfs.families import code_pattern, codes_in

#: Closed vocabulary the stance classifier answers from — anything else
#: (a refusal, a hedge, garbage) counts as "unclear" (Resolutions).
STANCES = ("support", "oppose", "modify", "unclear")

#: One docket's distinct-commenter count, with vs. without the family —
#: matches by literal code presence (``codes`` metadata) OR family
#: membership (``families`` metadata), the same two-way match
#: ``llm.evidence``'s code-cited/family-cited sources use, so a chunk
#: the indexer only stamped via family membership (no literal code
#: string in that particular chunk) still counts.
_DOCKET_COUNTS_SQL = text(
    "SELECT docket, "
    "count(DISTINCT item_key) FILTER (WHERE matched) AS n_items, "
    "count(DISTINCT item_key) AS n_total "
    "FROM ( "
    "SELECT e.cmetadata->>'docket' AS docket, e.cmetadata->>'item_key' AS item_key, "
    f"{code_or_family_predicate('e')} "
    "AS matched "
    "FROM langchain_pg_embedding e "
    "WHERE e.collection_id = (SELECT uuid FROM langchain_pg_collection WHERE name = :collection) "
    ") t "
    "WHERE docket IS NOT NULL AND docket <> '' "
    "GROUP BY docket ORDER BY docket"
)

#: The newest N distinct items' first matching chunk in one docket — the
#: same one-chunk-per-item windowing ``llm.evidence``'s family-cited SQL
#: uses (stage one: lowest ``seq`` per item_key; stage two: newest
#: ``date`` first), text truncated to 1,500 chars for the classifier
#: (Resolutions).
_STANCE_SAMPLE_SQL = text(
    "SELECT item_key, LEFT(document, 1500) FROM ( "
    "SELECT e.cmetadata->>'item_key' AS item_key, e.document AS document, "
    "e.cmetadata->>'date' AS date, "
    "row_number() OVER ( "
    "PARTITION BY e.cmetadata->>'item_key' "
    "ORDER BY COALESCE(NULLIF(e.cmetadata->>'seq', '')::int, 0) ASC "
    ") AS item_rn "
    "FROM langchain_pg_embedding e "
    "WHERE e.collection_id = (SELECT uuid FROM langchain_pg_collection WHERE name = :collection) "
    "AND e.cmetadata->>'docket' = :docket "
    "AND ( "
    f"{code_or_family_predicate('e')} "
    ") "
    ") one_per_item "
    "WHERE item_rn = 1 "
    "ORDER BY date DESC NULLS LAST "
    "LIMIT :n"
)

_COMMENTS_COLLECTION = "comments"


def docket_counts(
    engine: Any, family_key: str, codes: Sequence[str]
) -> list[tuple[str, int, int]]:
    """``(docket, n_items_with_family, n_items_total)`` for every docket
    that has any chunks in the ``comments`` collection — one round trip,
    ``GROUP BY docket``."""
    upper = [c.upper() for c in codes]
    with engine.begin() as conn:
        rows = conn.execute(
            _DOCKET_COUNTS_SQL,
            {
                "codes": upper,
                "family_key": family_key.upper(),
                "collection": _COMMENTS_COLLECTION,
            },
        ).fetchall()
    return [
        (docket, int(n_items or 0), int(n_total or 0))
        for docket, n_items, n_total in rows
    ]


def stance_sample(
    engine: Any, family_key: str, codes: Sequence[str], docket: str, *, n: int
) -> list[tuple[str, str]]:
    """``(item_key, text)`` for the newest *n* distinct commenters in
    *docket* who mentioned the family, first matching chunk, newest
    first, text capped to 1,500 chars."""
    if n <= 0:
        return []
    upper = [c.upper() for c in codes]
    with engine.begin() as conn:
        rows = conn.execute(
            _STANCE_SAMPLE_SQL,
            {
                "codes": upper,
                "family_key": family_key.upper(),
                "collection": _COMMENTS_COLLECTION,
                "docket": docket,
                "n": n,
            },
        ).fetchall()
    return [(item_key, text or "") for item_key, text in rows]


def _prompt_text(family_key: str, codes: Sequence[str], chunk: str) -> str:
    return (
        f"Comment on Medicare payment for {family_key} "
        f"(codes {' '.join(codes)}):\n{chunk}"
    )


def _stance_of(text: str, classify: Callable[[str, Sequence[str]], str | None]) -> str:
    stance = classify(text, STANCES)
    return stance if stance in STANCES else "unclear"


def classify_stances(
    sample: Sequence[tuple[str, str]],
    classify: Callable[[str, Sequence[str]], str | None],
) -> Counter:
    """Tally ``STANCES`` for a stance sample — an answer outside the
    closed vocabulary counts as ``"unclear"`` (Resolutions)."""
    return Counter(_stance_of(text, classify) for _, text in sample)


#: "Comment:" / "Response:" paragraph openers — a docket-comment excerpt
#: reproduced verbatim inside an FR rule's own preamble, always at the
#: very start of the paragraph. Tightened to no leading whitespace and no
#: space before the colon (case-insensitive, as before) so these agree
#: exactly with the ``LIKE 'Comment:%'``/``LIKE 'Response:%'`` SQL
#: prefilter in ``fr_pairs`` — a looser Python regex than the SQL LIKE
#: would count paragraphs the SQL never even fetched.
_COMMENT_RE = re.compile(r"^Comment:", re.I)
_RESPONSE_RE = re.compile(r"^Response:", re.I)


def fr_pairs(store: Any, codes: Sequence[str]) -> list[tuple[str, int, int, int]]:
    """``(item_key, rule_year, n_pairs, n_comment_paras)`` per rule item —
    the pre-2017 proxy for public reaction, since regulations.gov comment
    text isn't indexed that far back. ``n_pairs`` is the count of
    qualifying ``Comment:``/``Response:`` pairs naming one of *codes*
    (the family's share); ``n_comment_paras`` is the item's TOTAL count
    of ``Comment:`` paragraphs — every one in the rule, family or not,
    paired or not — the pool the family's share is drawn out of (the
    same n_items-out-of-n_total shape the docket rows use).

    Ruling A12: a ``Comment:`` paragraph pairs with the *next*
    ``Response:`` paragraph in the same item that comes before the next
    ``Comment:`` paragraph — intervening continuation paragraphs (the
    comment's own prose, wrapped across more than one ``fr_anchors``
    row) are allowed between them, and the code match runs over the
    whole span (comment + continuations + response), not just the
    comment/response paragraphs themselves. A ``Comment:`` with no
    ``Response:`` before the next ``Comment:`` (or the end of the item)
    is unpaired and doesn't count toward ``n_pairs`` — it still counts
    toward ``n_comment_paras``, which is unconditional. Rules with zero
    qualifying pairs are omitted (``n_comment_paras`` alone never keeps a
    rule in the output).

    One SQL per family, scoped to ``Comment:``/``Response:`` openers —
    unconditionally, every item's, which is what makes ``n_comment_paras``
    a true item-wide total even though the query is family-scoped — plus
    any paragraph naming one of *codes* (the same LIKE-prefilter idiom
    ``pfs.guidance.harvest`` uses) — safe under the pairing rule above: a
    dropped paragraph is prose that names no code and opens neither a
    comment nor a response, so it can neither create nor hide a pair.
    ``_COMMENT_RE``/``_RESPONSE_RE`` are anchored exactly like the SQL's
    ``LIKE 'Comment:%'``/``LIKE 'Response:%'`` (no leading whitespace, no
    space before the colon) so the Python pass never counts a paragraph
    the SQL wouldn't have fetched as an opener in the first place.
    """
    if not codes:
        return []
    con = store._con()  # noqa: SLF001 — same pattern as pfs.descriptors/pfs.guidance
    code_likes = " OR ".join("a.text LIKE ?" for _ in codes)
    params = [f"%{c.upper()}%" for c in codes]
    rows = con.execute(
        "SELECT a.item_key, a.p_id, a.text, i.title, i.date_published "
        "FROM fr_anchors a JOIN items i ON i.key = a.item_key "
        f"WHERE a.text LIKE 'Comment:%' OR a.text LIKE 'Response:%' OR ({code_likes}) "
        "ORDER BY a.item_key, a.p_id",
        params,
    ).fetchall()

    pattern = code_pattern(codes)
    by_item: dict[str, list[str]] = {}
    meta: dict[str, tuple[str, str]] = {}
    for item_key, _p_id, para_text, title, date_published in rows:
        by_item.setdefault(item_key, []).append(para_text or "")
        meta[item_key] = (title, date_published)

    out: list[tuple[str, int, int, int]] = []
    for item_key, paras in by_item.items():
        n_pairs = 0
        span: list[str] | None = None  # None: no comment currently open
        for para_text in paras:
            if _COMMENT_RE.match(para_text):
                # A new Comment: abandons any still-open, unresolved one
                # (Ruling A12: unpaired, never counted).
                span = [para_text]
                continue
            if _RESPONSE_RE.match(para_text):
                if span is not None:
                    span.append(para_text)
                    if codes_in(" ".join(span), pattern):
                        n_pairs += 1
                    span = None
                continue
            if span is not None:
                span.append(para_text)  # continuation paragraph
        if n_pairs:
            n_comment_paras = sum(1 for p in paras if _COMMENT_RE.match(p))
            title, date_published = meta[item_key]
            out.append(
                (
                    item_key,
                    rule_year_of(title, date_published),
                    n_pairs,
                    n_comment_paras,
                )
            )
    return sorted(out, key=lambda r: (r[1], r[0]))


def series(
    engine: Any,
    store: Any,
    family: str,
    codes: Sequence[str],
    *,
    classify: Callable[[str, Sequence[str]], str | None] | None = None,
    stance_sample_n: int = 0,
    name: str | None = None,
) -> list[Any]:
    """Every ``ReactionRow`` for one family (or ``--code`` pseudo-family):
    a docket row per docket with any comment chunks, then an fr-pairs row
    per rule with qualifying Comment/Response pairs. Stance columns stay
    zero (and ``sample_json`` ``"[]"``) unless *classify* is given and
    *stance_sample_n* > 0. *name* is the family's display name for the
    stance prompt (``FAMILIES[key].name`` for a registered family); it
    defaults to *family* itself, which is already the right label for a
    ``--code`` pseudo-family (there is no separate display name for a
    single code)."""
    from pfs.codetables import ReactionRow

    label = name or family
    docket_years = {d.id: d.comment_end_date[:4] for d in store.dockets()}

    rows: list[ReactionRow] = []
    for docket, n_items, n_total in docket_counts(engine, family, codes):
        year_str = docket_years.get(docket, "")
        year = int(year_str) if year_str.isdigit() else 0
        stance_counts: Counter = Counter()
        sample_json = "[]"
        if classify is not None and stance_sample_n > 0:
            sample = stance_sample(engine, family, codes, docket, n=stance_sample_n)
            stances = [
                (item_key, _stance_of(_prompt_text(label, codes, chunk), classify))
                for item_key, chunk in sample
            ]
            stance_counts = Counter(s for _, s in stances)
            sample_json = json.dumps([{"item_key": k, "stance": s} for k, s in stances])
        rows.append(
            ReactionRow(
                family=family,
                period=docket,
                period_kind="docket",
                year=year,
                n_items=n_items,
                n_total=n_total,
                stance_support=stance_counts.get("support", 0),
                stance_oppose=stance_counts.get("oppose", 0),
                stance_modify=stance_counts.get("modify", 0),
                stance_unclear=stance_counts.get("unclear", 0),
                sample_json=sample_json,
            )
        )

    for item_key, rule_year, n_pairs, n_comment_paras in fr_pairs(store, codes):
        rows.append(
            ReactionRow(
                family=family,
                period=item_key,
                period_kind="fr-pairs",
                year=rule_year,
                n_items=n_pairs,
                n_total=n_comment_paras,
                stance_support=0,
                stance_oppose=0,
                stance_modify=0,
                stance_unclear=0,
                sample_json="[]",
            )
        )
    return rows