stack/tests/pfs/test_payment.py

"""Tests for pfs.calcs.payment — payment(), direct_pe_cost(),
equipment_cost_per_use(), time_intensity_ratio(), and mppr_reduction().

All tests use polars DataFrames through @nw.narwhalify so the results
are plain polars DataFrames that can be inspected directly.
"""

from __future__ import annotations

import polars as pl
import pytest

from pfs.calcs.payment import (
    direct_pe_cost,
    equipment_cost_per_use,
    mppr_reduction,
    payment,
    time_intensity_ratio,
)
from tests.conftest import TEST_CF

# ── payment() ─────────────────────────────────────────────────────────────────


class TestPayment:
    """payment() joins RVU + GPCI and computes the fee-schedule amount.

    Formula:
        Payment = (Work_RVU × Work_GPCI
                 + PE_RVU   × PE_GPCI
                 + MP_RVU   × MP_GPCI) × Conversion_Factor
    """

    # ── return shape ──────────────────────────────────────────────────────────

    def test_returns_dataframe(self, rvu_df, gpci_df) -> None:
        result = payment(rvu_df, gpci_df, cf=TEST_CF)
        assert isinstance(result, pl.DataFrame)

    def test_row_count_matches_rvu(self, rvu_df, gpci_df) -> None:
        result = payment(rvu_df, gpci_df, cf=TEST_CF)
        assert len(result) == len(rvu_df)

    def test_payment_amount_column_present(self, rvu_df, gpci_df) -> None:
        result = payment(rvu_df, gpci_df, cf=TEST_CF)
        assert "payment_amount" in result.columns

    def test_original_columns_preserved(self, rvu_df, gpci_df) -> None:
        result = payment(rvu_df, gpci_df, cf=TEST_CF)
        for col in ("hcpcs", "work_rvu", "non_fac_pe_rvu", "mp_rvu"):
            assert col in result.columns, f"Column {col!r} missing from result"

    # ── non-facility calculation ──────────────────────────────────────────────

    def test_non_facility_payment_formula(self, rvu_df, gpci_df) -> None:
        """Non-facility payment matches the formula exactly for 99213."""
        result = payment(rvu_df, gpci_df, cf=TEST_CF, facility=False)
        row = result.filter(pl.col("hcpcs") == "99213")
        assert len(row) == 1

        # Expected: (0.97×1.0 + 1.04×0.998 + 0.07×0.633) × 32.7442
        expected = ((0.97 * 1.0) + (1.04 * 0.998) + (0.07 * 0.633)) * 32.7442
        actual = row["payment_amount"][0]
        assert actual == pytest.approx(expected, rel=1e-4)

    def test_non_facility_default(self, rvu_df, gpci_df) -> None:
        """facility=False is the default — results must match explicit False."""
        default_result = payment(rvu_df, gpci_df, cf=TEST_CF)
        explicit_result = payment(rvu_df, gpci_df, cf=TEST_CF, facility=False)
        assert default_result["payment_amount"].to_list() == pytest.approx(
            explicit_result["payment_amount"].to_list(), rel=1e-6
        )

    def test_non_facility_all_codes_positive(self, rvu_df, gpci_df) -> None:
        result = payment(rvu_df, gpci_df, cf=TEST_CF, facility=False)
        for val in result["payment_amount"].to_list():
            assert val > 0, f"payment_amount should be positive, got {val}"

    def test_non_facility_higher_rvu_higher_payment(self, rvu_df, gpci_df) -> None:
        """99215 has higher RVUs than 99213 — it must have a higher payment."""
        result = payment(rvu_df, gpci_df, cf=TEST_CF, facility=False)
        pay_99213 = result.filter(pl.col("hcpcs") == "99213")["payment_amount"][0]
        pay_99215 = result.filter(pl.col("hcpcs") == "99215")["payment_amount"][0]
        assert pay_99215 > pay_99213

    def test_non_facility_monotone_with_rvu(self, rvu_df, gpci_df) -> None:
        """Payment must be monotonically increasing with work_rvu across E&M codes."""
        result = payment(rvu_df, gpci_df, cf=TEST_CF, facility=False).sort("work_rvu")
        amounts = result["payment_amount"].to_list()
        for i in range(len(amounts) - 1):
            assert amounts[i] <= amounts[i + 1]

    # ── facility calculation ──────────────────────────────────────────────────

    def test_facility_payment_formula(self, rvu_df, gpci_df) -> None:
        """Facility payment uses fac_pe_rvu instead of non_fac_pe_rvu."""
        result = payment(rvu_df, gpci_df, cf=TEST_CF, facility=True)
        row = result.filter(pl.col("hcpcs") == "99213")

        expected = ((0.97 * 1.0) + (0.41 * 0.998) + (0.07 * 0.633)) * 32.7442
        actual = row["payment_amount"][0]
        assert actual == pytest.approx(expected, rel=1e-4)

    def test_facility_lower_than_non_facility(self, rvu_df, gpci_df) -> None:
        """Facility PE RVUs are always lower → facility payment < non-facility."""
        nf = payment(rvu_df, gpci_df, cf=TEST_CF, facility=False)[
            "payment_amount"
        ].to_list()
        f = payment(rvu_df, gpci_df, cf=TEST_CF, facility=True)[
            "payment_amount"
        ].to_list()
        for nf_pay, f_pay in zip(nf, f):
            assert f_pay < nf_pay, (
                f"Facility payment {f_pay} should be less than non-facility {nf_pay}"
            )

    def test_facility_positive(self, rvu_df, gpci_df) -> None:
        result = payment(rvu_df, gpci_df, cf=TEST_CF, facility=True)
        for val in result["payment_amount"].to_list():
            assert val > 0

    # ── join behavior ─────────────────────────────────────────────────────────

    def test_locality_join_works(self, rvu_df, gpci_df) -> None:
        """All rows join correctly when locality matches."""
        result = payment(rvu_df, gpci_df, cf=TEST_CF)
        # No nulls in payment_amount if join is a left join that matched
        null_count = result["payment_amount"].is_null().sum()
        assert null_count == 0

    def test_unmatched_locality_produces_null_payment(self) -> None:
        """A locality in RVU that has no matching GPCI → payment is null."""
        rvu = pl.DataFrame(
            {
                "hcpcs": ["99213"],
                "work_rvu": [0.97],
                "non_fac_pe_rvu": [1.04],
                "fac_pe_rvu": [0.41],
                "mp_rvu": [0.07],
                "mac": ["10212"],
                "locality": ["9999"],  # no match in gpci
            }
        )
        gpci = pl.DataFrame(
            {
                "mac": ["10212"],
                "locality": ["0201"],
                "work_gpci": [1.0],
                "pe_gpci": [0.998],
                "mp_gpci": [0.633],
            }
        )
        result = payment(rvu, gpci, cf=TEST_CF)
        # Left join — unmatched row should have null payment
        assert result["payment_amount"].is_null().sum() == 1

    def test_multiple_localities(self) -> None:
        """RVU rows for two localities should each use the correct GPCI."""
        rvu = pl.DataFrame(
            {
                "hcpcs": ["99213", "99213"],
                "work_rvu": [0.97, 0.97],
                "non_fac_pe_rvu": [1.04, 1.04],
                "fac_pe_rvu": [0.41, 0.41],
                "mp_rvu": [0.07, 0.07],
                "mac": ["10212", "10212"],
                "locality": ["0201", "0101"],
            }
        )
        gpci = pl.DataFrame(
            {
                "mac": ["10212", "10212"],
                "locality": ["0201", "0101"],
                "work_gpci": [1.0, 1.05],
                "pe_gpci": [0.998, 1.10],
                "mp_gpci": [0.633, 0.700],
            }
        )
        result = payment(rvu, gpci, cf=TEST_CF)
        assert len(result) == 2
        pay_0201 = result.filter(pl.col("locality") == "0201")["payment_amount"][0]
        pay_0101 = result.filter(pl.col("locality") == "0101")["payment_amount"][0]
        # 0101 has higher GPCIs → higher payment
        assert pay_0101 > pay_0201

    # ── conversion factor ─────────────────────────────────────────────────────

    def test_higher_conversion_factor_yields_higher_payment(self) -> None:
        """Doubling the conversion factor doubles the payment."""
        rvu = pl.DataFrame(
            {
                "hcpcs": ["99213"],
                "work_rvu": [0.97],
                "non_fac_pe_rvu": [1.04],
                "fac_pe_rvu": [0.41],
                "mp_rvu": [0.07],
                "mac": ["10212"],
                "locality": ["0201"],
            }
        )
        gpci = pl.DataFrame(
            {
                "mac": ["10212"],
                "locality": ["0201"],
                "work_gpci": [1.0],
                "pe_gpci": [0.998],
                "mp_gpci": [0.633],
            }
        )
        base_pay = payment(rvu, gpci, cf=TEST_CF)["payment_amount"][0]
        double_pay = payment(rvu, gpci, cf=TEST_CF * 2)["payment_amount"][0]
        assert double_pay == pytest.approx(base_pay * 2, rel=1e-6)

    # ── edge cases ────────────────────────────────────────────────────────────

    def test_zero_rvu_yields_zero_payment(self) -> None:
        rvu = pl.DataFrame(
            {
                "hcpcs": ["00000"],
                "work_rvu": [0.0],
                "non_fac_pe_rvu": [0.0],
                "fac_pe_rvu": [0.0],
                "mp_rvu": [0.0],
                "mac": ["10212"],
                "locality": ["0201"],
            }
        )
        gpci = pl.DataFrame(
            {
                "mac": ["10212"],
                "locality": ["0201"],
                "work_gpci": [1.0],
                "pe_gpci": [0.998],
                "mp_gpci": [0.633],
            }
        )
        result = payment(rvu, gpci, cf=TEST_CF)
        assert result["payment_amount"][0] == pytest.approx(0.0, abs=1e-10)

    def test_empty_rvu_returns_empty(self) -> None:
        rvu = pl.DataFrame(
            {
                "hcpcs": pl.Series([], dtype=pl.Utf8),
                "work_rvu": pl.Series([], dtype=pl.Float64),
                "non_fac_pe_rvu": pl.Series([], dtype=pl.Float64),
                "fac_pe_rvu": pl.Series([], dtype=pl.Float64),
                "mp_rvu": pl.Series([], dtype=pl.Float64),
                "mac": pl.Series([], dtype=pl.Utf8),
                "locality": pl.Series([], dtype=pl.Utf8),
            }
        )
        gpci = pl.DataFrame(
            {
                "mac": ["10212"],
                "locality": ["0201"],
                "work_gpci": [1.0],
                "pe_gpci": [0.998],
                "mp_gpci": [0.633],
            }
        )
        result = payment(rvu, gpci, cf=TEST_CF)
        assert len(result) == 0


# ── equipment_cost_per_use() ──────────────────────────────────────────────────


class TestEquipmentCostPerUse:
    """equipment_cost_per_use adds cost_per_minute = price / life / minutes_per_year."""

    # ── return shape ──────────────────────────────────────────────────────────

    def test_returns_dataframe(self, equipment_df) -> None:
        result = equipment_cost_per_use(equipment_df)
        assert isinstance(result, pl.DataFrame)

    def test_adds_cost_per_minute_column(self, equipment_df) -> None:
        result = equipment_cost_per_use(equipment_df)
        assert "cost_per_minute" in result.columns

    def test_row_count_unchanged(self, equipment_df) -> None:
        result = equipment_cost_per_use(equipment_df)
        assert len(result) == len(equipment_df)

    def test_original_columns_preserved(self, equipment_df) -> None:
        result = equipment_cost_per_use(equipment_df)
        for col in ("hcpcs", "unit_price", "useful_life", "minutes_per_year"):
            assert col in result.columns

    # ── formula without maintenance ───────────────────────────────────────────

    def test_formula_without_maintenance(self, equipment_df) -> None:
        """cost_per_minute = unit_price / useful_life / minutes_per_year."""
        result = equipment_cost_per_use(equipment_df)
        row = result.filter(pl.col("hcpcs") == "99213")
        expected = 1200.0 / 7.0 / 525600.0
        actual = row["cost_per_minute"][0]
        assert actual == pytest.approx(expected, rel=1e-6)

    def test_cost_per_minute_positive(self, equipment_df) -> None:
        result = equipment_cost_per_use(equipment_df)
        for val in result["cost_per_minute"].to_list():
            assert val > 0

    def test_higher_price_higher_cost(self) -> None:
        """Doubling unit_price doubles cost_per_minute."""
        eq_base = pl.DataFrame(
            {
                "hcpcs": ["99213"],
                "nf_minutes": [10.0],
                "f_minutes": [6.0],
                "unit_price": [1200.0],
                "useful_life": [7.0],
                "minutes_per_year": [525600.0],
            }
        )
        eq_double = eq_base.with_columns(pl.col("unit_price") * 2)

        base_cost = equipment_cost_per_use(eq_base)["cost_per_minute"][0]
        double_cost = equipment_cost_per_use(eq_double)["cost_per_minute"][0]
        assert double_cost == pytest.approx(base_cost * 2, rel=1e-6)

    def test_longer_life_lower_cost(self) -> None:
        """Longer useful life → lower cost per minute."""
        eq_short = pl.DataFrame(
            {
                "hcpcs": ["A"],
                "nf_minutes": [10.0],
                "f_minutes": [6.0],
                "unit_price": [1200.0],
                "useful_life": [5.0],
                "minutes_per_year": [525600.0],
            }
        )
        eq_long = eq_short.with_columns(pl.col("useful_life") * 2)

        short_cost = equipment_cost_per_use(eq_short)["cost_per_minute"][0]
        long_cost = equipment_cost_per_use(eq_long)["cost_per_minute"][0]
        assert long_cost < short_cost

    def test_more_minutes_per_year_lower_cost(self) -> None:
        """More annual minutes → lower per-minute cost (more utilization)."""
        eq = pl.DataFrame(
            {
                "hcpcs": ["A"],
                "nf_minutes": [10.0],
                "f_minutes": [6.0],
                "unit_price": [1200.0],
                "useful_life": [7.0],
                "minutes_per_year": [262800.0],  # half utilization
            }
        )
        eq_full = eq.with_columns(pl.col("minutes_per_year") * 2)

        half_util = equipment_cost_per_use(eq)["cost_per_minute"][0]
        full_util = equipment_cost_per_use(eq_full)["cost_per_minute"][0]
        assert full_util < half_util

    # ── formula with maintenance factor ──────────────────────────────────────

    def test_formula_with_maintenance_factor(
        self, equipment_with_maintenance_df
    ) -> None:
        """cost_per_minute = (price / life / mpy) × (1 + maintenance_factor)."""
        result = equipment_cost_per_use(equipment_with_maintenance_df)
        base = 1200.0 / 7.0 / 525600.0
        expected = base * (1 + 0.05)
        actual = result["cost_per_minute"][0]
        assert actual == pytest.approx(expected, rel=1e-6)

    def test_maintenance_increases_cost(
        self, equipment_df, equipment_with_maintenance_df
    ) -> None:
        """Adding maintenance factor must increase cost_per_minute."""
        base_cost = equipment_cost_per_use(equipment_df).filter(
            pl.col("hcpcs") == "99213"
        )["cost_per_minute"][0]
        maint_cost = equipment_cost_per_use(equipment_with_maintenance_df)[
            "cost_per_minute"
        ][0]
        assert maint_cost > base_cost

    def test_null_maintenance_factor_treated_as_zero(self) -> None:
        """A null maintenance_factor should behave identically to zero."""
        eq_null = pl.DataFrame(
            {
                "hcpcs": ["99213"],
                "nf_minutes": [10.0],
                "f_minutes": [6.0],
                "unit_price": [1200.0],
                "useful_life": [7.0],
                "minutes_per_year": [525600.0],
                "maintenance_factor": pl.Series([None], dtype=pl.Float64),
            }
        )
        eq_zero = pl.DataFrame(
            {
                "hcpcs": ["99213"],
                "nf_minutes": [10.0],
                "f_minutes": [6.0],
                "unit_price": [1200.0],
                "useful_life": [7.0],
                "minutes_per_year": [525600.0],
                "maintenance_factor": [0.0],
            }
        )
        null_cost = equipment_cost_per_use(eq_null)["cost_per_minute"][0]
        zero_cost = equipment_cost_per_use(eq_zero)["cost_per_minute"][0]
        assert null_cost == pytest.approx(zero_cost, rel=1e-9)

    def test_multiple_rows_all_have_cost(self, equipment_df) -> None:
        result = equipment_cost_per_use(equipment_df)
        assert result["cost_per_minute"].is_null().sum() == 0


# ── direct_pe_cost() ──────────────────────────────────────────────────────────


class TestDirectPeCost:
    """direct_pe_cost aggregates labor, supply, and equipment costs per HCPCS."""

    # ── return shape ──────────────────────────────────────────────────────────

    def test_returns_dataframe(self, labor_df, supply_df, equipment_df) -> None:
        result = direct_pe_cost(labor_df, supply_df, equipment_df)
        assert isinstance(result, pl.DataFrame)

    def test_contains_hcpcs_column(self, labor_df, supply_df, equipment_df) -> None:
        result = direct_pe_cost(labor_df, supply_df, equipment_df)
        assert "hcpcs" in result.columns

    def test_contains_nf_labor_cost(self, labor_df, supply_df, equipment_df) -> None:
        result = direct_pe_cost(labor_df, supply_df, equipment_df)
        assert "nf_labor_cost" in result.columns

    def test_contains_f_labor_cost(self, labor_df, supply_df, equipment_df) -> None:
        result = direct_pe_cost(labor_df, supply_df, equipment_df)
        assert "f_labor_cost" in result.columns

    def test_contains_nf_supply_cost(self, labor_df, supply_df, equipment_df) -> None:
        result = direct_pe_cost(labor_df, supply_df, equipment_df)
        assert "nf_supply_cost" in result.columns

    def test_contains_f_supply_cost(self, labor_df, supply_df, equipment_df) -> None:
        result = direct_pe_cost(labor_df, supply_df, equipment_df)
        assert "f_supply_cost" in result.columns

    def test_contains_nf_equipment_cost(
        self, labor_df, supply_df, equipment_df
    ) -> None:
        result = direct_pe_cost(labor_df, supply_df, equipment_df)
        assert "nf_equipment_cost" in result.columns

    def test_contains_f_equipment_cost(self, labor_df, supply_df, equipment_df) -> None:
        result = direct_pe_cost(labor_df, supply_df, equipment_df)
        assert "f_equipment_cost" in result.columns

    def test_contains_nf_direct_pe(self, labor_df, supply_df, equipment_df) -> None:
        result = direct_pe_cost(labor_df, supply_df, equipment_df)
        assert "nf_direct_pe" in result.columns

    def test_contains_f_direct_pe(self, labor_df, supply_df, equipment_df) -> None:
        result = direct_pe_cost(labor_df, supply_df, equipment_df)
        assert "f_direct_pe" in result.columns

    # ── labor aggregation ─────────────────────────────────────────────────────

    def test_labor_cost_aggregated_per_hcpcs(
        self, labor_df, supply_df, equipment_df
    ) -> None:
        """99213 has two labor rows — they should be summed."""
        result = direct_pe_cost(labor_df, supply_df, equipment_df)
        row_99213 = result.filter(pl.col("hcpcs") == "99213")
        assert len(row_99213) == 1  # one output row per HCPCS

    def test_nf_labor_cost_99213(self, labor_df, supply_df, equipment_df) -> None:
        """99213 nf_labor = (14×0.63) + (5×0.49) = 8.82 + 2.45 = 11.27."""
        result = direct_pe_cost(labor_df, supply_df, equipment_df)
        val = result.filter(pl.col("hcpcs") == "99213")["nf_labor_cost"][0]
        expected = (14.0 * 0.63) + (5.0 * 0.49)
        assert val == pytest.approx(expected, rel=1e-4)

    def test_f_labor_cost_99213(self, labor_df, supply_df, equipment_df) -> None:
        """99213 f_labor = (9×0.63) + (3×0.49) = 5.67 + 1.47 = 7.14."""
        result = direct_pe_cost(labor_df, supply_df, equipment_df)
        val = result.filter(pl.col("hcpcs") == "99213")["f_labor_cost"][0]
        expected = (9.0 * 0.63) + (3.0 * 0.49)
        assert val == pytest.approx(expected, rel=1e-4)

    # ── supply aggregation ────────────────────────────────────────────────────

    def test_nf_supply_cost_99213(self, labor_df, supply_df, equipment_df) -> None:
        """99213 nf_supply = 2 × 0.12 = 0.24."""
        result = direct_pe_cost(labor_df, supply_df, equipment_df)
        val = result.filter(pl.col("hcpcs") == "99213")["nf_supply_cost"][0]
        expected = 2.0 * 0.12
        assert val == pytest.approx(expected, rel=1e-6)

    def test_f_supply_cost_99213(self, labor_df, supply_df, equipment_df) -> None:
        """99213 f_supply = 1 × 0.12 = 0.12."""
        result = direct_pe_cost(labor_df, supply_df, equipment_df)
        val = result.filter(pl.col("hcpcs") == "99213")["f_supply_cost"][0]
        expected = 1.0 * 0.12
        assert val == pytest.approx(expected, rel=1e-6)

    # ── direct PE totals ──────────────────────────────────────────────────────

    def test_nf_direct_pe_is_sum_of_components(
        self, labor_df, supply_df, equipment_df
    ) -> None:
        """nf_direct_pe = nf_labor + nf_supply + nf_equipment (nulls as 0)."""
        result = direct_pe_cost(labor_df, supply_df, equipment_df)
        row = result.filter(pl.col("hcpcs") == "99213")
        nf_pe = row["nf_direct_pe"][0]
        nf_labor = row["nf_labor_cost"][0] or 0.0
        nf_supply = row["nf_supply_cost"][0] or 0.0
        nf_equip = row["nf_equipment_cost"][0] or 0.0
        assert nf_pe == pytest.approx(nf_labor + nf_supply + nf_equip, rel=1e-5)

    def test_f_direct_pe_is_sum_of_components(
        self, labor_df, supply_df, equipment_df
    ) -> None:
        """f_direct_pe = f_labor + f_supply + f_equipment."""
        result = direct_pe_cost(labor_df, supply_df, equipment_df)
        row = result.filter(pl.col("hcpcs") == "99213")
        f_pe = row["f_direct_pe"][0]
        f_labor = row["f_labor_cost"][0] or 0.0
        f_supply = row["f_supply_cost"][0] or 0.0
        f_equip = row["f_equipment_cost"][0] or 0.0
        assert f_pe == pytest.approx(f_labor + f_supply + f_equip, rel=1e-5)

    def test_direct_pe_nonnegative(self, labor_df, supply_df, equipment_df) -> None:
        result = direct_pe_cost(labor_df, supply_df, equipment_df)
        for col in ("nf_direct_pe", "f_direct_pe"):
            for val in result[col].to_list():
                if val is not None:
                    assert val >= 0, f"{col} has negative value: {val}"

    def test_hcpcs_with_no_supply_still_aggregated(
        self, labor_df, equipment_df
    ) -> None:
        """A HCPCS code present in labor but absent from supply should still appear."""
        supply_empty = pl.DataFrame(
            {
                "hcpcs": pl.Series([], dtype=pl.Utf8),
                "nf_quantity": pl.Series([], dtype=pl.Float64),
                "f_quantity": pl.Series([], dtype=pl.Float64),
                "unit_price": pl.Series([], dtype=pl.Float64),
            }
        )
        result = direct_pe_cost(labor_df, supply_empty, equipment_df)
        assert isinstance(result, pl.DataFrame)
        assert len(result) > 0


# ── time_intensity_ratio() ────────────────────────────────────────────────────


class TestTimeIntensityRatio:
    """time_intensity_ratio computes work_rvu / time for each HCPCS code."""

    # ── return shape ──────────────────────────────────────────────────────────

    def test_returns_dataframe(self, work_time_df) -> None:
        result = time_intensity_ratio(work_time_df)
        assert isinstance(result, pl.DataFrame)

    def test_row_count_preserved(self, work_time_df) -> None:
        result = time_intensity_ratio(work_time_df)
        assert len(result) == len(work_time_df)

    def test_adds_intensity_total(self, work_time_df) -> None:
        result = time_intensity_ratio(work_time_df)
        assert "intensity_total" in result.columns

    def test_adds_intensity_intra(self, work_time_df) -> None:
        result = time_intensity_ratio(work_time_df)
        assert "intensity_intra" in result.columns

    def test_original_columns_preserved(self, work_time_df) -> None:
        result = time_intensity_ratio(work_time_df)
        for col in ("hcpcs", "work_rvu", "total_time", "intra_service_time"):
            assert col in result.columns

    # ── formula ───────────────────────────────────────────────────────────────

    def test_intensity_total_formula(self, work_time_df) -> None:
        """intensity_total = work_rvu / total_time for each row."""
        result = time_intensity_ratio(work_time_df)
        for i in range(len(work_time_df)):
            wrvu = work_time_df["work_rvu"][i]
            total = work_time_df["total_time"][i]
            expected = wrvu / total
            actual = result["intensity_total"][i]
            assert actual == pytest.approx(expected, rel=1e-6), (
                f"Row {i}: intensity_total mismatch"
            )

    def test_intensity_intra_formula(self, work_time_df) -> None:
        """intensity_intra = work_rvu / intra_service_time for each row."""
        result = time_intensity_ratio(work_time_df)
        for i in range(len(work_time_df)):
            wrvu = work_time_df["work_rvu"][i]
            intra = work_time_df["intra_service_time"][i]
            expected = wrvu / intra
            actual = result["intensity_intra"][i]
            assert actual == pytest.approx(expected, rel=1e-6), (
                f"Row {i}: intensity_intra mismatch"
            )

    def test_intensity_intra_higher_than_total(self, work_time_df) -> None:
        """intra_service_time < total_time → intensity_intra > intensity_total."""
        result = time_intensity_ratio(work_time_df)
        for i in range(len(result)):
            intra = result["intensity_intra"][i]
            total = result["intensity_total"][i]
            assert intra > total, (
                f"Row {i}: intensity_intra ({intra}) should exceed intensity_total ({total})"
            )

    def test_higher_rvu_higher_intensity(self, work_time_df) -> None:
        """99214 (1.50 RVU / 31 min = 0.0484) > 99213 (0.97 / 21 = 0.0462).

        Note: 99215 (2.11/46=0.0459) is actually LOWER than 99213 because
        its time grows faster than its RVU — intensity is not monotone with
        code level. Use 99214 vs 99213 as the reliable comparison pair.
        """
        result = time_intensity_ratio(work_time_df)
        i_99213 = result.filter(pl.col("hcpcs") == "99213")["intensity_total"][0]
        i_99214 = result.filter(pl.col("hcpcs") == "99214")["intensity_total"][0]
        assert i_99214 > i_99213

    # ── without reference code ────────────────────────────────────────────────

    def test_no_ratio_columns_without_reference(self, work_time_df) -> None:
        """Without reference_hcpcs, ratio_total and ratio_intra should NOT be added."""
        result = time_intensity_ratio(work_time_df)
        assert "ratio_total" not in result.columns
        assert "ratio_intra" not in result.columns

    # ── with reference code ───────────────────────────────────────────────────

    def test_ratio_columns_added_with_reference(self, work_time_df) -> None:
        result = time_intensity_ratio(work_time_df, reference_hcpcs="99213")
        assert "ratio_total" in result.columns
        assert "ratio_intra" in result.columns

    def test_reference_code_ratio_is_one(self, work_time_df) -> None:
        """The reference code's ratio to itself is always 1.0."""
        result = time_intensity_ratio(work_time_df, reference_hcpcs="99213")
        ref_ratio = result.filter(pl.col("hcpcs") == "99213")["ratio_total"][0]
        assert ref_ratio == pytest.approx(1.0, rel=1e-6)

    def test_reference_code_ratio_intra_is_one(self, work_time_df) -> None:
        result = time_intensity_ratio(work_time_df, reference_hcpcs="99213")
        ref_ratio = result.filter(pl.col("hcpcs") == "99213")["ratio_intra"][0]
        assert ref_ratio == pytest.approx(1.0, rel=1e-6)

    def test_higher_code_ratio_above_one(self, work_time_df) -> None:
        """99214 has higher intensity than 99213 → ratio > 1 when 99213 is reference.

        99215 has lower intensity than 99213 (time grows faster than RVU), so
        use 99214 as the code that reliably exceeds the 99213 reference.
        """
        result = time_intensity_ratio(work_time_df, reference_hcpcs="99213")
        ratio_99214 = result.filter(pl.col("hcpcs") == "99214")["ratio_total"][0]
        assert ratio_99214 > 1.0

    def test_ratio_formula_total(self, work_time_df) -> None:
        """ratio_total = intensity_total / ref_intensity_total."""
        result = time_intensity_ratio(work_time_df, reference_hcpcs="99213")
        ref_intensity = result.filter(pl.col("hcpcs") == "99213")["intensity_total"][0]
        for i in range(len(result)):
            expected = result["intensity_total"][i] / ref_intensity
            actual = result["ratio_total"][i]
            assert actual == pytest.approx(expected, rel=1e-6)

    def test_invalid_reference_code_raises_or_returns_null(self, work_time_df) -> None:
        """A reference HCPCS that does not exist in the DataFrame should not silently
        produce wrong results — either raise or produce null ratios."""
        try:
            result = time_intensity_ratio(work_time_df, reference_hcpcs="XXXXX")
            # If it doesn't raise, ratio columns may be null or NaN — that's acceptable
            if "ratio_total" in result.columns:
                pass  # Implementation handles missing ref gracefully
        except Exception:
            pass  # Raising is also acceptable


# ── mppr_reduction() ──────────────────────────────────────────────────────────


class TestMpprReduction:
    """mppr_reduction applies 50% PE reduction to 2nd+ procedures in a session."""

    # ── return shape ──────────────────────────────────────────────────────────

    def test_returns_dataframe(self, mppr_claims_df) -> None:
        result = mppr_reduction(mppr_claims_df)
        assert isinstance(result, pl.DataFrame)

    def test_row_count_preserved(self, mppr_claims_df) -> None:
        result = mppr_reduction(mppr_claims_df)
        assert len(result) == len(mppr_claims_df)

    def test_adds_pe_rvu_adjusted(self, mppr_claims_df) -> None:
        result = mppr_reduction(mppr_claims_df)
        assert "pe_rvu_adjusted" in result.columns

    def test_adds_mppr_rank(self, mppr_claims_df) -> None:
        result = mppr_reduction(mppr_claims_df)
        assert "mppr_rank" in result.columns

    def test_original_columns_preserved(self, mppr_claims_df) -> None:
        result = mppr_reduction(mppr_claims_df)
        for col in ("claim_id", "hcpcs", "session_id", "non_fac_pe_rvu"):
            assert col in result.columns

    # ── ranking ───────────────────────────────────────────────────────────────

    def test_highest_pe_gets_rank_one(self, mppr_claims_df) -> None:
        """The procedure with the highest PE RVU in the session gets rank 1."""
        result = mppr_reduction(mppr_claims_df)
        # 99214 has non_fac_pe_rvu=1.56, 99213 has 1.04 — 99214 should be rank 1
        row_99214 = result.filter(pl.col("hcpcs") == "99214")
        assert row_99214["mppr_rank"][0] == 1

    def test_lower_pe_gets_rank_two(self, mppr_claims_df) -> None:
        result = mppr_reduction(mppr_claims_df)
        row_99213 = result.filter(pl.col("hcpcs") == "99213")
        assert row_99213["mppr_rank"][0] == 2

    def test_rank_one_gets_full_payment(self, mppr_claims_df) -> None:
        """Rank-1 procedure keeps its full PE RVU."""
        result = mppr_reduction(mppr_claims_df)
        rank1 = result.filter(pl.col("mppr_rank") == 1)
        original_pe = rank1["non_fac_pe_rvu"][0]
        adjusted_pe = rank1["pe_rvu_adjusted"][0]
        assert adjusted_pe == pytest.approx(original_pe, rel=1e-6)

    def test_rank_two_gets_50pct_reduction_by_default(self, mppr_claims_df) -> None:
        """Rank-2+ procedure gets 50% of its PE RVU by default."""
        result = mppr_reduction(mppr_claims_df)
        rank2 = result.filter(pl.col("mppr_rank") == 2)
        original_pe = rank2["non_fac_pe_rvu"][0]
        adjusted_pe = rank2["pe_rvu_adjusted"][0]
        assert adjusted_pe == pytest.approx(original_pe * 0.50, rel=1e-6)

    # ── custom reduction percentage ───────────────────────────────────────────

    def test_custom_reduction_pct(self, mppr_claims_df) -> None:
        """A custom reduction_pct applies correctly to rank-2+ procedures."""
        result = mppr_reduction(mppr_claims_df, reduction_pct=0.25)
        rank2 = result.filter(pl.col("mppr_rank") == 2)
        original_pe = rank2["non_fac_pe_rvu"][0]
        adjusted_pe = rank2["pe_rvu_adjusted"][0]
        assert adjusted_pe == pytest.approx(original_pe * 0.75, rel=1e-6)

    def test_zero_reduction_pct_no_reduction(self, mppr_claims_df) -> None:
        """reduction_pct=0 means all procedures get full PE payment."""
        result = mppr_reduction(mppr_claims_df, reduction_pct=0.0)
        for i in range(len(result)):
            original = result["non_fac_pe_rvu"][i]
            adjusted = result["pe_rvu_adjusted"][i]
            assert adjusted == pytest.approx(original, rel=1e-6)

    def test_full_reduction_pct_zeroes_rank2(self, mppr_claims_df) -> None:
        """reduction_pct=1.0 zeros out PE for rank-2+ procedures."""
        result = mppr_reduction(mppr_claims_df, reduction_pct=1.0)
        rank2 = result.filter(pl.col("mppr_rank") == 2)
        assert rank2["pe_rvu_adjusted"][0] == pytest.approx(0.0, abs=1e-9)

    # ── custom pe_component ───────────────────────────────────────────────────

    def test_facility_pe_component(self) -> None:
        """mppr_reduction can target fac_pe_rvu instead of non_fac_pe_rvu."""
        claims = pl.DataFrame(
            {
                "claim_id": ["CLM001", "CLM001"],
                "hcpcs": ["99214", "99213"],
                "session_id": ["SES001", "SES001"],
                "non_fac_pe_rvu": [1.56, 1.04],
                "fac_pe_rvu": [0.63, 0.41],
            }
        )
        result = mppr_reduction(claims, pe_component="fac_pe_rvu")
        rank1 = result.filter(pl.col("mppr_rank") == 1)
        rank2 = result.filter(pl.col("mppr_rank") == 2)
        # Rank 1 should be 99214 (fac_pe=0.63 > 99213 fac_pe=0.41)
        assert rank1["fac_pe_rvu"][0] == pytest.approx(rank1["pe_rvu_adjusted"][0])
        assert rank2["pe_rvu_adjusted"][0] == pytest.approx(
            rank2["fac_pe_rvu"][0] * 0.50, rel=1e-6
        )

    # ── single procedure session ──────────────────────────────────────────────

    def test_single_procedure_gets_rank_one(self) -> None:
        """A session with only one procedure should get rank 1 and full PE."""
        single = pl.DataFrame(
            {
                "claim_id": ["CLM001"],
                "hcpcs": ["99213"],
                "session_id": ["SES001"],
                "non_fac_pe_rvu": [1.04],
            }
        )
        result = mppr_reduction(single)
        assert result["mppr_rank"][0] == 1
        assert result["pe_rvu_adjusted"][0] == pytest.approx(1.04, rel=1e-6)

    # ── multiple sessions ─────────────────────────────────────────────────────

    def test_rank_is_session_scoped(self) -> None:
        """Each session has its own independent ranking."""
        claims = pl.DataFrame(
            {
                "claim_id": ["A", "A", "B", "B"],
                "hcpcs": ["99214", "99213", "99215", "99213"],
                "session_id": ["SES001", "SES001", "SES002", "SES002"],
                "non_fac_pe_rvu": [1.56, 1.04, 2.22, 1.04],
            }
        )
        result = mppr_reduction(claims)
        # SES001: 99214 is rank 1, 99213 is rank 2
        # SES002: 99215 is rank 1, 99213 is rank 2
        ses1_rank1 = result.filter(
            (pl.col("session_id") == "SES001") & (pl.col("mppr_rank") == 1)
        )
        ses2_rank1 = result.filter(
            (pl.col("session_id") == "SES002") & (pl.col("mppr_rank") == 1)
        )
        assert ses1_rank1["hcpcs"][0] == "99214"
        assert ses2_rank1["hcpcs"][0] == "99215"

    def test_adjusted_pe_always_nonnegative(self, mppr_claims_df) -> None:
        result = mppr_reduction(mppr_claims_df)
        for val in result["pe_rvu_adjusted"].to_list():
            if val is not None:
                assert val >= 0