stack/tests/pfs/test_cms_summaries.py

"""pfs.cms_summaries — CMS's own Comment:/Response: summaries (dossier task 7)."""

from __future__ import annotations

import json
from pathlib import Path
from types import SimpleNamespace

import duckdb
import pytest

from bib.item import Rule
from bib.store import Store
from pfs.cms_summaries import (
    SUMMARY_PROMPT_VERSION,
    SUMMARY_STANCES,
    quantifier_of,
    response_action_of,
    summaries,
    summary_stance,
)
from pfs.codetables import (
    CmsSummaryRow,
    ensure_tables,
    read_cms_summaries,
    write_cms_summaries,
)
from pfs.reaction import PairSpan, fr_pair_spans
from pfs.runlog import RunLog

_RAW = json.loads(
    (Path(__file__).parent / "fixtures" / "fr_comment_pairs.json").read_text()
)
# Live Federal Register paragraph text, keyed by paragraph id per rule item.
CY2023 = {int(k.split(":")[1]): v for k, v in _RAW.items() if k.startswith("PETM")}
CY2008 = {int(k.split(":")[1]): v for k, v in _RAW.items() if k.startswith("54GT")}
# R44/R45 live paragraphs of other rules (XU44ZSI9 CY2006, YYAWUFQX CY2007,
# DE2VH9PD CY2015, JCWARTJD CY2027)
OTHER = {
    int(k.split(":")[1]): v
    for k, v in _RAW.items()
    if not k.startswith(("PETM", "54GT", "X-"))
}
TEXT_X = {k[2:]: v for k, v in _RAW.items() if k.startswith("X-")}
TEXT = {**CY2023, **CY2008, **OTHER}

FAM = SimpleNamespace(
    key="HOME-VISITS",
    codes=("99341", "99342", "99350", "99326"),
    pair_phrases=("home or residence services",),
)


@pytest.mark.parametrize(
    "p,q",
    [(1679, "majority"), (1681, "few"), (1674, "unspecified"), (2241, "unspecified")],
)
def test_quantifier(p, q):
    assert quantifier_of(TEXT[p]) == q


@pytest.mark.parametrize(
    "text,q",
    [
        ("Comment: All commenters agreed.", "all"),
        ("Comment: All of the commenters agreed.", "all"),
        ("Comment: Most commenters agreed.", "majority"),
        ("Comment: Many commenters agreed.", "many"),
        ("Comment: Many of the commenters agreed.", "many"),
        # R44: a quantifier must modify "commenters"
        ("Comment: We received many comments and several documents.", "unspecified"),
        ("Comment: Some commenters asked for several settings.", "some"),
        ("Comment: A few of the commenters agreed.", "few"),
        ("Comment: Numerous commenters agreed.", "many"),
        ("Comment: Several commenters agreed.", "several"),
        ("Comment: Some commenters agreed.", "some"),
        ("Comment: One commenter agreed.", "one"),
        ("Comment: Another commenter agreed.", "one"),
        ("Comment: A commenter agreed.", "one"),
        ("Comment: Commenters in somewhere else.", "unspecified"),
        ("Comment: Manyfold commenters.", "unspecified"),
    ],
)
def test_quantifier_words(text, q):
    assert quantifier_of(text) == q


def test_cy2027_1214_some_commenters_not_several_settings():
    # live CY2027 text: "Some commenters ... across several settings"
    assert quantifier_of(TEXT[1214]) == "some"


@pytest.mark.parametrize("p", [759, 1736, 1347])
def test_response_action_plain_refusals_in_live_text(p):
    # R45: live responses that plainly decline (no authority / inappropriate /
    # not warranted); every other unclear response stays unclear
    assert response_action_of(TEXT[p]) == "not_finalized"


def test_a_few_comments_is_few_not_one():
    # live CY2023 ¶3620: "A few comments ... One commenter indicated"
    assert quantifier_of(TEXT_X["PETM225D:3620"]) == "few"


@pytest.mark.parametrize("p", [1231, 1165, 1670])
def test_a_response_that_does_more_than_thank_is_not_acknowledged(p):
    # live: "We are not revising this policy" / "we are clarifying" / "we are adding"
    assert response_action_of(TEXT[p]) != "acknowledged"


def test_thanks_with_a_but_clause_is_not_acknowledged():
    # live PETM225D ¶406: "We appreciate ..., but we would need code-level PE details"
    assert response_action_of(TEXT_X["PETM225D:406"]) == "unclear"
    assert response_action_of(TEXT_X["YBM4IZUS:429"]) == "acknowledged"


def test_pure_thanks_is_acknowledged():
    assert response_action_of(TEXT[1680]) == "acknowledged"
    assert (
        response_action_of("Response: We thank commenters. We appreciate it.")
        == "acknowledged"
    )


def test_conflicting_cues_are_unclear():
    # a finalized-as-proposed cue next to a not-finalized cue is unclear
    assert (
        response_action_of(
            "Response: We are not finalizing X. We are finalizing Y as proposed."
        )
        == "unclear"
    )


def test_not_believe_finalizing_is_still_a_finalization():
    # R46a, live CY2027 ¶5899: "we do not believe that finalizing, as proposed, ..."
    assert response_action_of(TEXT[5899]) == "finalize_as_proposed"


def test_finalizing_beyond_a_comma_clause_is_as_proposed():
    # R47, live 9V89NVZG ¶1624: "as proposed" more than 120 characters after "finalizing"
    assert response_action_of(TEXT[1624]) == "finalize_as_proposed"


@pytest.mark.parametrize("p", [864, 3377])
def test_revis_far_from_finalizing_is_not_finalize_modified(p):
    # R47a, live DE2VH9PD ¶864 ("revised codes") and WHNGCYHB ¶3377 ("the revision")
    assert response_action_of(TEXT[p]) != "finalize_modified"


def test_thanks_with_an_action_verb_is_not_acknowledged():
    # R48, live K292HMSA ¶3866: "We appreciate the commenters' support and are finalizing ..."
    assert response_action_of(TEXT[3866]) == "unclear"


def test_inappropriate_needs_a_refusal_object():
    # live YYAWUFQX ¶1476: "it would be inappropriate to have collected" is not a refusal
    assert response_action_of(TEXT[1476]) != "not_finalized"
    assert response_action_of(TEXT[1581]) != "not_finalized"
    # R46a, live WHNGCYHB ¶3310: "inappropriate to apply the adjustment" is not a refusal
    assert response_action_of(TEXT[3310]) != "not_finalized"
    assert (
        response_action_of(
            "Response: We believe it would be inappropriate to finalize this."
        )
        == "not_finalized"
    )


def test_a_few_is_not_one():
    assert quantifier_of("A few commenters said") == "few"


@pytest.mark.parametrize(
    "p,a",
    [
        (1675, "finalize_as_proposed"),
        (1680, "acknowledged"),
        (1682, "finalize_as_proposed"),
        (2242, "finalize_modified"),
    ],
)
def test_response_action(p, a):
    assert response_action_of(TEXT[p]) == a


@pytest.mark.parametrize(
    "text,a",
    [
        ("Response: We are finalizing it with modifications.", "finalize_modified"),
        (
            "Response: We are finalizing the value, revised from 2.0.",
            "finalize_modified",
        ),
        ("Response: We are not finalizing this as proposed.", "not_finalized"),
        ("Response: We decline to adopt it.", "not_finalized"),
        ("Response: We are not adopting that.", "not_finalized"),
        ("Response: Thank you for the comment.", "acknowledged"),
        ("Response: We disagree.", "unclear"),
    ],
)
def test_response_action_words(text, a):
    assert response_action_of(text) == a


def _store():
    s = Store(":memory:")
    con = s._con()  # noqa: SLF001
    keys = {}
    for tag, title, date, texts in (
        ("c23", "Medicare Program; CY 2023 PFS Final Rule", "2022-11-18", CY2023),
        ("c08", "Medicare Program; CY 2008 PFS Final Rule", "2007-12-27", CY2008),
    ):
        key = s.create(Rule(title=title, date_published=date))
        keys[tag] = key
        con.executemany(
            "INSERT INTO fr_anchors (item_key, p_id, page, ordinal, text) VALUES (?,?,?,?,?)",
            [(key, p, 1, i, t) for i, (p, t) in enumerate(sorted(texts.items()))],
        )
    con.commit()
    return s, keys


def test_fr_pair_spans_by_code_and_phrase():
    s, keys = _store()
    try:
        by_code = fr_pair_spans(s, ("99341",))
        assert [(x.comment_p_id, x.response_p_id) for x in by_code] == [
            (1681, 1682),
            (1683, 1684),
        ]
        sp = by_code[0]
        assert isinstance(sp, PairSpan) and sp.rule_year == 2023
        assert sp.item_key == keys["c23"] and sp.codes[0] == "99341"
        assert sp.comment_text == TEXT[1681] and sp.response_text == TEXT[1682]
        both = fr_pair_spans(s, ("99341",), ("Home or Residence Services",))
        assert [x.comment_p_id for x in both] == [1674, 1679, 1681, 1683]
        assert fr_pair_spans(s, ("99341",), ()) == by_code  # no phrases: code only
        phrase_only = [x for x in both if x.comment_p_id == 1679][0]
        assert phrase_only.codes == ()
        assert fr_pair_spans(s, (), ()) == []
    finally:
        s.close()


def test_fr_pair_spans_phrase_only_input():
    s, _ = _store()
    try:
        out = fr_pair_spans(s, (), ("home or residence services",))
        assert [x.comment_p_id for x in out] == [1674, 1679, 1681]
    finally:
        s.close()


def test_summaries_pair_and_classify(tmp_path):
    s, keys = _store()
    seen = []

    def classify(text, choices):
        seen.append((text, tuple(choices)))
        return "support" if "majority" in text.lower() else "modify"

    try:
        rows = summaries(
            s, FAM, classify=classify, model="m", log=RunLog(tmp_path / "l.jsonl")
        )
        by = {(r.item_key, r.comment_p_id): r for r in rows}
        r = by[(keys["c23"], 1679)]
        assert (r.response_p_id, r.quantifier, r.stance) == (
            1680,
            "majority",
            "support",
        )
        assert r.family == "HOME-VISITS" and r.rule_year == 2023
        assert r.rule_kind == "final" and r.codes == ""
        assert (
            r.response_action == "acknowledged"
            and r.extractor == f"model:m@{SUMMARY_PROMPT_VERSION}"
        )
        assert by[(keys["c23"], 1681)].codes == "99341,99342,99350"
        assert by[(keys["c08"], 2241)].codes == "99326,99350"
        assert by[(keys["c08"], 2241)].response_action == "finalize_modified"
        assert seen[0][0].startswith("CMS's summary of public comments:")
        assert seen[0][1] == SUMMARY_STANCES
    finally:
        s.close()


def test_summaries_without_classifier_are_regex_unclear():
    s, _ = _store()
    try:
        rows = summaries(s, FAM)
        assert rows and {r.stance for r in rows} == {"unclear"}
        assert {r.extractor for r in rows} == {"regex"}
    finally:
        s.close()


def test_summary_stance_outside_vocabulary_is_unclear():
    assert summary_stance("x", lambda t, c: "support") == "support"
    assert summary_stance("x", lambda t, c: "banana") == "unclear"
    assert summary_stance("x", lambda t, c: None) == "unclear"


def test_summaries_ledger_resume_and_replay(tmp_path):
    s, keys = _store()
    log = RunLog(tmp_path / "l.jsonl")
    calls = []

    def classify(text, choices):
        calls.append(text)
        return "oppose"

    try:
        counts: dict = {}
        first = summaries(s, FAM, classify=classify, model="m", log=log, counts=counts)
        n = len(first)
        assert counts["done"] == n and len(calls) == n
        assert any(f"{keys['c23']}:1679" == e.unit for e in log._load())  # noqa: SLF001
        again: dict = {}
        summaries(s, FAM, classify=classify, model="m", log=log, counts=again)
        assert again["skipped"] == n and len(calls) == n  # no new calls
        # --no-llm replay: same rows, zero calls
        replay = summaries(s, FAM, log=log, replay_model="m")
        assert replay == first and len(calls) == n
        # unrecorded under a different model: unclear / regex
        other = summaries(s, FAM, log=log, replay_model="other")
        assert {r.stance for r in other} == {"unclear"}
        assert {r.extractor for r in other} == {"regex"}
        # limit caps new calls; the rest stay regex
        fresh = RunLog(tmp_path / "l2.jsonl")
        capped = summaries(s, FAM, classify=classify, model="m", log=fresh, limit=1)
        assert sum(r.extractor != "regex" for r in capped) == 1
    finally:
        s.close()


def test_summaries_failed_unit_stays_regex(tmp_path):
    s, _ = _store()
    log = RunLog(tmp_path / "l.jsonl")
    try:
        # a logged failure (output None) is not replayed as a stance
        rows = summaries(s, FAM, classify=lambda t, c: "support", model="m", log=log)
        unit = f"{rows[0].item_key}:{rows[0].comment_p_id}"
        e = log.get(unit, *_hash_of(log, unit), "m", SUMMARY_PROMPT_VERSION)
        assert e.output == {"stance": "support"}
        log2 = RunLog(tmp_path / "l3.jsonl")
        log2.put(unit, e.input_hash, "m", SUMMARY_PROMPT_VERSION, None)
        out = summaries(s, FAM, log=log2, replay_model="m")
        assert out[0].stance == "unclear" and out[0].extractor == "regex"
    finally:
        s.close()


def _hash_of(log, unit):
    e = next(e for e in log._load() if e.unit == unit)  # noqa: SLF001
    return (e.input_hash,)


def test_summaries_skip_correction_items():
    s, _ = _store()
    try:
        corr = s.create(
            Rule(title="CY 2023 PFS Correction notice", date_published="2023-01-06")
        )
        con = s._con()  # noqa: SLF001
        con.execute(
            "INSERT INTO fr_anchors (item_key, p_id, page, ordinal, text) VALUES (?,?,?,?,?)",
            (corr, 1, 1, 1, "Comment: A commenter on 99341."),
        )
        con.execute(
            "INSERT INTO fr_anchors (item_key, p_id, page, ordinal, text) VALUES (?,?,?,?,?)",
            (corr, 2, 1, 2, "Response: We thank the commenter."),
        )
        con.commit()
        # not vacuous: the pair is found, but the correction item is skipped
        assert corr in {x.item_key for x in fr_pair_spans(s, FAM.codes)}
        assert corr not in {r.item_key for r in summaries(s, FAM)}
    finally:
        s.close()


def test_write_read_roundtrip():
    con = duckdb.connect(":memory:")
    ensure_tables(con)
    row = CmsSummaryRow(
        "HOME-VISITS",
        "K",
        2023,
        "final",
        1679,
        1680,
        "",
        "majority",
        "support",
        "acknowledged",
        "regex",
    )
    other = CmsSummaryRow(
        "OTHER", "K", 2023, "final", 1, 2, "99341", "one", "unclear", "unclear", "regex"
    )
    assert write_cms_summaries(con, "HOME-VISITS", [row]) == 1
    write_cms_summaries(con, "OTHER", [other])
    assert write_cms_summaries(con, "HOME-VISITS", [row]) == 1  # delete-then-insert
    assert read_cms_summaries(con, "HOME-VISITS") == [row]
    assert write_cms_summaries(con, "HOME-VISITS", []) == 0
    assert read_cms_summaries(con, "HOME-VISITS") == []
    assert read_cms_summaries(con, "OTHER") == [other]