stack/tests/pfs/test_reaction.py

"""pfs.reaction — public-reaction series per code family (#690, task 7)."""

from __future__ import annotations

import dataclasses
import json
from unittest.mock import MagicMock

import duckdb
import pytest

from bib.dockets import Docket
from bib.item import Rule
from bib.store import Store
from pfs.codetables import ReactionRow, ensure_tables, read_reaction, write_reaction
from pfs.reaction import (
    STANCES,
    classify_stances,
    docket_counts,
    fr_pairs,
    series,
    stance_sample,
)

# ── fake pgvector engine ──────────────────────────────────────────────


def _engine(rows):
    engine = MagicMock()
    conn = engine.begin.return_value.__enter__.return_value
    conn.execute.return_value.fetchall.return_value = rows
    return engine, conn


def _multi_engine(docket_rows=None, stance_rows=None):
    """A fake engine whose rows depend on which of the two queries is
    running — distinguished by the ``docket``/``n`` bind params that
    only ``_STANCE_SAMPLE_SQL`` carries."""
    engine = MagicMock()
    conn = engine.begin.return_value.__enter__.return_value

    def _execute(_sql, params):
        result = MagicMock()
        if "docket" in params:
            result.fetchall.return_value = stance_rows or []
        else:
            result.fetchall.return_value = docket_rows or []
        return result

    conn.execute.side_effect = _execute
    return engine, conn


class TestDocketCounts:
    def test_maps_rows_and_binds_params(self):
        engine, conn = _engine([("CMS-2023-0121", 40, 200), ("CMS-2024-0453", 3, 150)])
        out = docket_counts(engine, "CCM", ("99490", "99439"))
        assert out == [("CMS-2023-0121", 40, 200), ("CMS-2024-0453", 3, 150)]
        params = conn.execute.call_args.args[1]
        assert params["codes"] == ["99490", "99439"]
        assert params["family_key"] == "CCM"
        assert params["collection"] == "comments"

    def test_null_counts_become_zero(self):
        engine, _conn = _engine([("CMS-2023-0121", None, 0)])
        assert docket_counts(engine, "CCM", ()) == [("CMS-2023-0121", 0, 0)]


class TestStanceSample:
    def test_maps_rows_and_binds_params(self):
        engine, conn = _engine([("C1", "chunk one"), ("C2", "chunk two")])
        out = stance_sample(engine, "CCM", ("99490",), "CMS-2023-0121", n=5)
        assert out == [("C1", "chunk one"), ("C2", "chunk two")]
        params = conn.execute.call_args.args[1]
        assert params["docket"] == "CMS-2023-0121"
        assert params["n"] == 5
        assert params["family_key"] == "CCM"

    def test_n_zero_short_circuits_without_a_query(self):
        engine, conn = _engine([("C1", "x")])
        assert stance_sample(engine, "CCM", ("99490",), "CMS-2023-0121", n=0) == []
        conn.execute.assert_not_called()


class TestClassifyStances:
    def test_tallies_known_slugs(self):
        sample = [("C1", "t1"), ("C2", "t2"), ("C3", "t3")]
        answers = {"t1": "support", "t2": "oppose", "t3": "support"}
        counts = classify_stances(sample, lambda text, choices: answers[text])
        assert counts["support"] == 2 and counts["oppose"] == 1

    def test_out_of_vocabulary_answer_counts_as_unclear(self):
        sample = [("C1", "t1")]
        counts = classify_stances(sample, lambda text, choices: "maybe-ish")
        assert counts["unclear"] == 1

    def test_none_answer_counts_as_unclear(self):
        sample = [("C1", "t1")]
        counts = classify_stances(sample, lambda text, choices: None)
        assert counts["unclear"] == 1

    def test_every_stance_is_a_known_slug(self):
        assert set(STANCES) == {"support", "oppose", "modify", "unclear"}


# ── fr_pairs ────────────────────────────────────────────────────────


@pytest.fixture
def store_with_pairs():
    s = Store(":memory:")
    item_key = s.create(Rule(title="CY2015 PFS final rule"))
    con = s._con()  # noqa: SLF001
    con.executemany(
        "INSERT INTO fr_anchors (item_key, p_id, page, ordinal, text) VALUES (?,?,?,?,?)",
        [
            (
                item_key,
                10,
                100,
                1,
                "Comment: Commenters asked why CPT code 99490 (chronic care "
                "management) is valued as proposed.",
            ),
            (
                item_key,
                11,
                100,
                2,
                "Response: We are finalizing the RVUs for 99490 as proposed.",
            ),
            (item_key, 12, 101, 3, "Comment: A commenter raised an unrelated point."),
            (item_key, 13, 101, 4, "Response: We disagree with the commenter."),
            # a Comment: with no immediately-following Response: (next para
            # is prose) must never be counted.
            (item_key, 14, 102, 5, "Comment: Another commenter mentioned 99490."),
            (item_key, 15, 102, 6, "We did not adopt this suggestion."),
        ],
    )
    con.commit()
    yield s, item_key
    s.close()


@pytest.fixture
def store_with_continuation_pair():
    """Ruling A12: a Comment: paragraph with no code of its own, a
    continuation paragraph (no Comment:/Response: prefix) naming a code,
    then the Response: — must still count as one pair."""
    s = Store(":memory:")
    item_key = s.create(Rule(title="CY2016 PFS final rule"))
    con = s._con()  # noqa: SLF001
    con.executemany(
        "INSERT INTO fr_anchors (item_key, p_id, page, ordinal, text) VALUES (?,?,?,?,?)",
        [
            (
                item_key,
                20,
                200,
                1,
                "Comment: Commenters raised concerns about staff time requirements.",
            ),
            (
                item_key,
                21,
                200,
                2,
                "Providers noted that 99490 requires clinical staff supervision "
                "that is hard to bill separately.",
            ),
            (
                item_key,
                22,
                200,
                3,
                "Response: We agree and are clarifying supervision requirements.",
            ),
        ],
    )
    con.commit()
    yield s, item_key
    s.close()


class TestFrPairs:
    def test_counts_only_pairs_naming_a_code(self, store_with_pairs):
        # n_total is the item's TOTAL Comment: paragraph count (family or
        # not) — store_with_pairs has three: ¶10, ¶12, ¶14.
        s, item_key = store_with_pairs
        out = fr_pairs(s, ("99490",))
        assert out == [(item_key, 2015, 1, 3)]

    def test_no_codes_returns_empty(self, store_with_pairs):
        s, _item_key = store_with_pairs
        assert fr_pairs(s, ()) == []

    def test_unmatched_code_yields_no_rows(self, store_with_pairs):
        s, _item_key = store_with_pairs
        assert fr_pairs(s, ("99437",)) == []

    def test_dangling_comment_without_response_is_not_counted(self, store_with_pairs):
        # ¶14's "Comment:" is followed by ¶15, which is not a "Response:"
        # paragraph — it must not inflate the count even though it names
        # 99490.
        s, item_key = store_with_pairs
        out = fr_pairs(s, ("99490",))
        assert out == [(item_key, 2015, 1, 3)]

    def test_code_named_only_in_a_continuation_paragraph_still_pairs(
        self, store_with_continuation_pair
    ):
        # Ruling A12: intervening continuation paragraphs are allowed
        # between Comment: and Response:, and the code match runs over
        # the whole span — neither the Comment: nor the Response:
        # paragraph names 99490 here, only the continuation between them.
        # This item has one Comment: paragraph total, so n_total is 1.
        s, item_key = store_with_continuation_pair
        out = fr_pairs(s, ("99490",))
        assert out == [(item_key, 2016, 1, 1)]

    def test_new_comment_abandons_an_unresolved_open_one(self):
        # A Comment: immediately followed by another Comment: (no
        # Response: in between) leaves the first one unpaired; the
        # second one, properly followed by a Response:, still counts.
        s = Store(":memory:")
        item_key = s.create(Rule(title="CY2017 PFS final rule"))
        con = s._con()  # noqa: SLF001
        con.executemany(
            "INSERT INTO fr_anchors (item_key, p_id, page, ordinal, text) "
            "VALUES (?,?,?,?,?)",
            [
                (item_key, 1, 1, 1, "Comment: A commenter mentioned 99490 in passing."),
                (item_key, 2, 1, 2, "Comment: A second, unrelated 99490 comment."),
                (item_key, 3, 1, 3, "Response: We are finalizing 99490 as proposed."),
            ],
        )
        con.commit()
        try:
            out = fr_pairs(s, ("99490",))
        finally:
            s.close()
        # Two Comment: paragraphs total in this item (¶1 abandoned, ¶2
        # paired), so n_total is 2 even though only one pair qualifies.
        assert out == [(item_key, 2017, 1, 2)]


# ── series ──────────────────────────────────────────────────────────


class TestSeries:
    @pytest.fixture
    def store(self, store_with_pairs):
        s, item_key = store_with_pairs
        s.docket_upsert(
            Docket(id="CMS-2023-0121", rule_cms_id="R1", comment_end_date="2023-06-01")
        )
        return s, item_key

    def test_docket_and_fr_pairs_rows_without_stance(self, store):
        s, item_key = store
        engine, _conn = _multi_engine(docket_rows=[("CMS-2023-0121", 5, 20)])
        rows = series(engine, s, "CCM", ("99490",))
        by_kind = {r.period_kind: r for r in rows}
        docket_row = by_kind["docket"]
        assert (
            docket_row.period,
            docket_row.year,
            docket_row.n_items,
            docket_row.n_total,
        ) == (
            "CMS-2023-0121",
            2023,
            5,
            20,
        )
        assert docket_row.stance_support == 0 and docket_row.sample_json == "[]"
        fr_row = by_kind["fr-pairs"]
        assert (fr_row.period, fr_row.year, fr_row.n_items, fr_row.n_total) == (
            item_key,
            2015,
            1,
            3,
        )

    def test_stance_sample_classifies_and_fills_sample_json(self, store):
        s, _item_key = store
        engine, _conn = _multi_engine(
            docket_rows=[("CMS-2023-0121", 5, 20)],
            stance_rows=[("C1", "we support this"), ("C2", "we oppose this")],
        )

        def fake_classify(text, choices):
            return "support" if "support" in text else "oppose"

        rows = series(
            engine, s, "CCM", ("99490",), classify=fake_classify, stance_sample_n=2
        )
        docket_row = next(r for r in rows if r.period_kind == "docket")
        assert docket_row.stance_support == 1 and docket_row.stance_oppose == 1
        sample = json.loads(docket_row.sample_json)
        assert {s_["item_key"] for s_ in sample} == {"C1", "C2"}
        assert {s_["stance"] for s_ in sample} == {"support", "oppose"}

    def test_name_overrides_family_key_in_the_stance_prompt(self, store):
        s, _item_key = store
        engine, _conn = _multi_engine(
            docket_rows=[("CMS-2023-0121", 5, 20)],
            stance_rows=[("C1", "a comment")],
        )
        seen_prompts = []

        def fake_classify(text, choices):
            seen_prompts.append(text)
            return "support"

        series(
            engine,
            s,
            "CCM",
            ("99490",),
            classify=fake_classify,
            stance_sample_n=1,
            name="Chronic Care Management",
        )
        assert len(seen_prompts) == 1
        assert "Chronic Care Management" in seen_prompts[0]
        assert "CCM" not in seen_prompts[0]

    def test_zero_stance_sample_skips_classification(self, store):
        s, _item_key = store
        engine, conn = _multi_engine(docket_rows=[("CMS-2023-0121", 5, 20)])
        rows = series(
            engine,
            s,
            "CCM",
            ("99490",),
            classify=lambda t, c: "support",
            stance_sample_n=0,
        )
        docket_row = next(r for r in rows if r.period_kind == "docket")
        assert docket_row.stance_support == 0
        # only the docket_counts query ran — no stance_sample round trip
        assert conn.execute.call_count == 1

    def test_code_pseudo_family(self, store):
        s, _item_key = store
        engine, conn = _multi_engine(docket_rows=[("CMS-2023-0121", 2, 20)])
        rows = series(engine, s, "G2211", ("G2211",))
        assert [r.family for r in rows if r.period_kind == "docket"] == ["G2211"]
        params = conn.execute.call_args.args[1]
        assert params["codes"] == ["G2211"]


# ── DDL / write / read round trip ────────────────────────────────────


@pytest.fixture
def con():
    c = duckdb.connect(":memory:")
    ensure_tables(c)
    yield c
    c.close()


class TestReactionTable:
    def _row(
        self, family="CCM", period="CMS-2023-0121", period_kind="docket", year=2023
    ):
        return ReactionRow(family, period, period_kind, year, 5, 20, 1, 0, 0, 1, "[]")

    def test_write_then_read_round_trips(self, con):
        rows = [self._row(), self._row(period="R1", period_kind="fr-pairs", year=2015)]
        n = write_reaction(con, "CCM", rows)
        assert n == 2
        got = read_reaction(con, "CCM")
        assert {(r.period, r.period_kind) for r in got} == {
            ("CMS-2023-0121", "docket"),
            ("R1", "fr-pairs"),
        }

    def test_write_is_delete_then_insert_per_family(self, con):
        write_reaction(con, "CCM", [self._row(period="CMS-2023-0121")])
        write_reaction(con, "CCM", [self._row(period="CMS-2024-0453")])
        got = read_reaction(con, "CCM")
        assert [r.period for r in got] == ["CMS-2024-0453"]

    def test_other_family_untouched(self, con):
        write_reaction(con, "G2211", [self._row(family="G2211")])
        write_reaction(con, "CCM", [])
        assert len(read_reaction(con, "G2211")) == 1


class TestEventsJson:
    def test_events_json_roundtrip_and_null_tolerated(self, con):
        ev = json.dumps(
            {"F:2023:other": {"support": 1, "modify": 0, "oppose": 0, "unclear": 0}}
        )
        row = ReactionRow(
            "CCM", "D1", "docket", 2023, 5, 20, 1, 0, 0, 0, "[]", ev, 4, 9
        )
        write_reaction(con, "CCM", [row])
        assert read_reaction(con, "CCM") == [row]
        never = ReactionRow("CCM", "D2", "docket", 2024, 1, 2, 0, 0, 0, 0, "[]")
        assert (never.n_classified, never.n_commenters) == (None, None)
        tagged = ReactionRow(
            "CCM",
            "D3",
            "docket",
            2025,
            1,
            2,
            0,
            0,
            0,
            0,
            "[]",
            "{}",
            3,
            4,
            "stance-v4:m",
            0.83,
        )
        write_reaction(con, "CCM", [row, never, tagged])
        assert read_reaction(con, "CCM")[2].stance_agreement == 0.83
        assert read_reaction(con, "CCM")[2].stance_method == "stance-v4:m"
        write_reaction(con, "CCM", [row, never])
        assert read_reaction(con, "CCM")[1].n_classified is None
        assert (
            ReactionRow("CCM", "D", "docket", 1, 0, 0, 0, 0, 0, 0, "[]").events_json
            == "{}"
        )
        # a row written before the column existed reads back "{}"
        con.execute(
            "INSERT INTO pfs.code_reaction (family, period, period_kind, year, n_items, "
            "n_total, stance_support, stance_oppose, stance_modify, stance_unclear, "
            "sample_json) VALUES ('OLD', 'D0', 'docket', 2020, 1, 2, 0, 0, 0, 0, '[]')"
        )
        assert read_reaction(con, "OLD")[0].events_json == "{}"
        assert read_reaction(con, "OLD")[0].n_commenters is None

    def test_ensure_tables_adds_the_column_to_an_old_table(self):
        c = duckdb.connect(":memory:")
        c.execute("CREATE SCHEMA pfs")
        c.execute(
            "CREATE TABLE pfs.code_reaction (family VARCHAR, period VARCHAR, "
            "period_kind VARCHAR, year INTEGER, n_items INTEGER, n_total INTEGER, "
            "stance_support INTEGER, stance_oppose INTEGER, stance_modify INTEGER, "
            "stance_unclear INTEGER, sample_json VARCHAR)"
        )
        c.execute(
            "INSERT INTO pfs.code_reaction VALUES ('A','D','docket',1,1,1,0,0,0,0,'[]')"
        )
        short = read_reaction(c, "A")[0]  # an old replica: rows without the columns
        assert (short.events_json, short.n_classified) == ("{}", None)
        assert (short.stance_method, short.stance_agreement) == (None, None)
        ensure_tables(c)
        got = read_reaction(c, "A")[0]
        assert (got.events_json, got.n_classified, got.n_commenters) == (
            "{}",
            None,
            None,
        )
        assert got.stance_eval is None

    def test_stance_eval_roundtrips_and_an_older_table_gains_it(self):
        """R66a: ``stance_eval`` (JSON) is a trailing column; a table from
        round 3 (``stance_agreement`` last) keeps its rows readable and gains
        the column in place."""
        c = duckdb.connect(":memory:")
        c.execute("CREATE SCHEMA pfs")
        c.execute(
            "CREATE TABLE pfs.code_reaction (family VARCHAR, period VARCHAR, "
            "period_kind VARCHAR, year INTEGER, n_items INTEGER, n_total INTEGER, "
            "stance_support INTEGER, stance_oppose INTEGER, stance_modify INTEGER, "
            "stance_unclear INTEGER, sample_json VARCHAR, events_json VARCHAR, "
            "n_classified INTEGER, n_commenters INTEGER, stance_method VARCHAR, "
            "stance_agreement DOUBLE)"
        )
        c.execute(
            "INSERT INTO pfs.code_reaction VALUES ('A','D','docket',1,1,1,0,0,0,0,"
            "'[]','{}',3,4,'stance-v4:m',0.83)"
        )
        old = read_reaction(c, "A")[0]
        assert (old.stance_method, old.stance_agreement, old.stance_eval) == (
            "stance-v4:m",
            0.83,
            None,
        )
        ensure_tables(c)
        assert read_reaction(c, "A")[0] == old
        ev = json.dumps({"cell": 0.83, "recall": 0.556})
        row = dataclasses.replace(
            old, period="D2", stance_agreement=None, stance_eval=ev
        )
        write_reaction(c, "A", [old, row])
        assert read_reaction(c, "A") == [old, row]


class TestFrPairSpansRefactor:
    def test_counts_derive_from_spans(self, store_with_pairs):
        from pfs.reaction import fr_pair_spans

        s, item_key = store_with_pairs
        spans = fr_pair_spans(s, ("99490",))
        assert [
            (x.item_key, x.rule_year, x.comment_p_id, x.response_p_id) for x in spans
        ] == [(item_key, 2015, 10, 11)]
        assert spans[0].codes == ("99490",)
        assert fr_pairs(s, ("99490",)) == [(item_key, 2015, len(spans), 3)]