stack/tests/aco/test_express_cms_quality_measures.py

"""Tests for aco.express.cms_quality_measures.

Covers UAMCC (NQF #2888), ACR (NQF #1789), and HWR express functions
implementing the CMS hospital admission and readmission quality measures.
"""

from __future__ import annotations

from datetime import date

import polars as pl
import pytest

from aco.express.cms_quality_measures import (
    _INJURY_ACCIDENT_CCS,
    _PROC_COMPLICATION_CCS,
    MCC_GROUPS,
    SPECIALTY_COHORTS,
    acr_int_index_admission,
    acr_int_planned_readmission,
    acr_int_specialty_cohort,
    acr_performance_period,
    hwr_int_denominator,
    hwr_int_planned_readmission,
    hwr_performance_period,
    stg_medical_claim,
    stg_medical_claim_condition,
    uamcc_int_denominator,
    uamcc_int_denominator_exclusion,
    uamcc_int_mcc_cohort,
    uamcc_int_outcome_exclusion,
    uamcc_int_person_time,
    uamcc_int_planned_admission,
    uamcc_performance_period,
)

# ═══════════════════════════════════════════════════════════════════════════════
# Shared fixtures
# ═══════════════════════════════════════════════════════════════════════════════


@pytest.fixture
def uamcc_period_df() -> pl.DataFrame:
    """UAMCC performance period anchor row."""
    return pl.DataFrame(
        {
            "measure_id": ["UAMCC"],
            "measure_name": ["All-Cause Unplanned Admissions for Patients with MCCs"],
            "nqf_id": ["2888"],
            "performance_year": [2025],
            "performance_period_begin": [date(2025, 1, 1)],
            "performance_period_end": [date(2025, 12, 31)],
            "lookback_period_begin": [date(2024, 1, 1)],
            "lookback_period_end": [date(2024, 12, 31)],
        }
    )


@pytest.fixture
def acr_period_df() -> pl.DataFrame:
    """ACR performance period anchor row."""
    return pl.DataFrame(
        {
            "measure_id": ["ACR"],
            "measure_name": ["Risk-Standardized, All-Condition Readmission"],
            "nqf_id": ["1789"],
            "performance_year": [2025],
            "performance_period_begin": [date(2025, 1, 1)],
            "performance_period_end": [date(2025, 12, 31)],
        }
    )


@pytest.fixture
def hwr_period_df() -> pl.DataFrame:
    """HWR performance period anchor row."""
    return pl.DataFrame(
        {
            "measure_id": ["HWR"],
            "measure_name": ["Hospital-wide, 30-Day, All-cause Unplanned Readmission"],
            "performance_year": [2025],
            "performance_period_begin": [date(2025, 1, 1)],
            "performance_period_end": [date(2025, 12, 31)],
        }
    )


@pytest.fixture
def patient_df() -> pl.DataFrame:
    """Patient table with varying birth and death dates."""
    return pl.DataFrame(
        {
            "person_id": ["P001", "P002", "P003", "P004", "P005"],
            "sex": ["M", "F", "M", "F", "M"],
            "birth_date": [
                date(1945, 6, 1),  # age ~79 at 2025-01-01
                date(1948, 3, 15),  # age ~76
                date(1965, 1, 1),  # age 60 — under 66, should be excluded
                date(1950, 7, 4),  # age ~74
                date(1940, 11, 20),  # age ~84
            ],
            "death_date": [None, None, None, date(2025, 6, 1), None],
        },
        schema={
            "person_id": pl.String,
            "sex": pl.String,
            "birth_date": pl.Date,
            "death_date": pl.Date,
        },
    )


@pytest.fixture
def uamcc_cohort_value_set_df() -> pl.DataFrame:
    """UAMCC Cohort value set — ICD-10 codes for each of the 9 MCC groups."""
    return pl.DataFrame(
        {
            "chronic_condition_group": [
                "AMI",
                "AMI",
                "HEART_FAILURE",
                "HEART_FAILURE",
                "DIABETES",
                "DIABETES",
                "CKD",
                "COPD_ASTHMA",
                "AFIB",
                "ALZHEIMER",
                "DEPRESSION",
                "STROKE_TIA",
            ],
            "icd_10_cm": [
                "I21.0",
                "I21.1",
                "I50.1",
                "I50.2",
                "E11.9",
                "E10.9",
                "N18.3",
                "J44.1",
                "I48.0",
                "G30.0",
                "F32.0",
                "I63.5",
            ],
            "label": [
                "STEMI anterior",
                "STEMI inferior",
                "Systolic HF",
                "Diastolic HF",
                "Type 2 DM uncomplicated",
                "Type 1 DM uncomplicated",
                "CKD stage 3",
                "COPD exacerbation",
                "AFib paroxysmal",
                "Alzheimer's",
                "Major depressive episode",
                "Cerebral infarction",
            ],
            "claims_to_qualify": ["1 inpatient OR 2 outpatient"] * 12,
            "diagnoses_used": ["Principal or secondary"] * 12,
            "lookback_years": [1] * 12,
            "additional_notes": [None] * 12,
        }
    )


@pytest.fixture
def stg_medical_claim_df() -> pl.DataFrame:
    """Staged medical claims with principal diagnosis (rank-1 condition)."""
    return pl.DataFrame(
        {
            "claim_id": [
                "CLM001",
                "CLM002",
                "CLM003",
                "CLM004",
                "CLM005",
                "CLM006",
                "CLM007",
                "CLM008",
            ],
            "person_id": [
                "P001",
                "P001",
                "P002",
                "P002",
                "P004",
                "P004",
                "P005",
                "P005",
            ],
            "claim_start_date": [
                date(2024, 3, 1),
                date(2024, 5, 1),
                date(2024, 2, 1),
                date(2024, 4, 1),
                date(2024, 1, 1),
                date(2024, 6, 1),
                date(2024, 3, 1),
                date(2024, 7, 1),
            ],
            "claim_end_date": [
                date(2024, 3, 3),
                date(2024, 5, 3),
                date(2024, 2, 2),
                date(2024, 4, 2),
                date(2024, 1, 2),
                date(2024, 6, 3),
                date(2024, 3, 4),
                date(2024, 7, 3),
            ],
            "principal_diagnosis_code": [
                "I21.0",
                "I50.1",
                "E11.9",
                "N18.3",
                "J44.1",
                "I48.0",
                "G30.0",
                "F32.0",
            ],
            "hcpcs_code": pl.Series([None] * 8, dtype=pl.String),
            "place_of_service_code": ["21"] * 8,
        }
    )


@pytest.fixture
def stg_medical_claim_condition_df() -> pl.DataFrame:
    """Staged claim-condition pairs (all diagnosis positions)."""
    return pl.DataFrame(
        {
            "claim_id": [
                "CLM001",
                "CLM002",
                "CLM003",
                "CLM004",
                "CLM005",
                "CLM006",
                "CLM007",
                "CLM008",
            ],
            "person_id": [
                "P001",
                "P001",
                "P002",
                "P002",
                "P004",
                "P004",
                "P005",
                "P005",
            ],
            "claim_start_date": [
                date(2024, 3, 1),
                date(2024, 5, 1),
                date(2024, 2, 1),
                date(2024, 4, 1),
                date(2024, 1, 1),
                date(2024, 6, 1),
                date(2024, 3, 1),
                date(2024, 7, 1),
            ],
            "normalized_code": [
                "I21.0",  # P001 AMI
                "I50.1",  # P001 Heart failure → 2 conditions
                "E11.9",  # P002 Diabetes
                "N18.3",  # P002 CKD → 2 conditions
                "J44.1",  # P004 COPD
                "I48.0",  # P004 AFib → 2 conditions
                "G30.0",  # P005 Alzheimer
                "F32.0",  # P005 Depression → 2 conditions
            ],
        }
    )


@pytest.fixture
def encounter_df() -> pl.DataFrame:
    """Core encounters including acute inpatient stays."""
    return pl.DataFrame(
        {
            "encounter_id": ["ENC001", "ENC002", "ENC003", "ENC004", "ENC005"],
            "person_id": ["P001", "P001", "P002", "P004", "P005"],
            "encounter_type": [
                "acute inpatient",
                "acute inpatient",
                "acute inpatient",
                "acute inpatient",
                "outpatient",
            ],
            "encounter_start_date": [
                date(2025, 2, 1),
                date(2025, 4, 5),
                date(2025, 3, 1),
                date(2025, 1, 15),
                date(2025, 5, 1),
            ],
            "encounter_end_date": [
                date(2025, 2, 5),
                date(2025, 4, 8),
                date(2025, 3, 4),
                date(2025, 1, 18),
                date(2025, 5, 1),
            ],
            "length_of_stay": [4, 3, 3, 3, 0],
            "discharge_disposition_code": ["01", "01", "01", "01", None],
            "facility_id": ["H001", "H001", "H002", "H003", None],
            "primary_diagnosis_code": ["I50.1", "I21.0", "E11.9", "J44.1", "Z00.0"],
            "ccs_diagnosis_category": ["108", "100", "49", "127", None],
            "drg_code_type": ["MS-DRG"] * 4 + [None],
            "drg_code": ["291", "282", "637", "190", None],
            "encounter_group": ["claims"] * 5,
        },
        schema={
            "encounter_id": pl.String,
            "person_id": pl.String,
            "encounter_type": pl.String,
            "encounter_start_date": pl.Date,
            "encounter_end_date": pl.Date,
            "length_of_stay": pl.Int32,
            "discharge_disposition_code": pl.String,
            "facility_id": pl.String,
            "primary_diagnosis_code": pl.String,
            "ccs_diagnosis_category": pl.String,
            "drg_code_type": pl.String,
            "drg_code": pl.String,
            "encounter_group": pl.String,
        },
    )


def _make_paa1_df() -> pl.DataFrame:
    return pl.DataFrame(
        {
            "ccs_procedure_category": ["64", "105", "116", "142", "171"],
            "description": [
                "Bone marrow transplant",
                "Kidney transplant",
                "Heart transplant",
                "Partial excision bone",
                "Amputation lower extremity",
            ],
        }
    )


def _make_paa2_df() -> pl.DataFrame:
    return pl.DataFrame(
        {
            "ccs_diagnosis_category": ["45", "254"],
            "description": [
                "Maintenance chemotherapy; radiotherapy",
                "Rehabilitation care",
            ],
        }
    )


def _make_paa3_df() -> pl.DataFrame:
    return pl.DataFrame(
        {
            "code_type": ["CCS", "CCS", "ICD-10-PCS"],
            "category_or_code": ["1", "2", "0DTJ4ZZ"],
            "description": ["Incision CNS", "Endoscopy", "Resection appendix"],
            "associated_ccs_category": [None, None, "74"],
        }
    )


def _make_paa4_df() -> pl.DataFrame:
    return pl.DataFrame(
        {
            "code_type": ["CCS", "CCS", "ICD-10-CM"],
            "category_or_code": ["2", "100", "I21.0"],
            "description": ["Septicemia", "Acute MI", "STEMI anterior"],
            "associated_ccs_category": [None, None, "100"],
        }
    )


def _make_ccs_icd10_cm_df() -> pl.DataFrame:
    """Small CCS → ICD-10-CM crosswalk for testing."""
    return pl.DataFrame(
        {
            "icd_10_cm": [
                "I50.1",
                "I21.0",
                "E11.9",
                "N18.3",
                "J44.1",
                "I48.0",
                "G30.0",
                "F32.0",
                "I63.5",
                "T82.7XXA",  # complication of device
                "W19.XXXA",  # fall
                "Z51.11",  # chemo maintenance → CCS 45
            ],
            "description": [
                "Systolic HF",
                "STEMI",
                "T2DM",
                "CKD3",
                "COPD exac",
                "AFib",
                "Alzheimer",
                "MDD",
                "Stroke",
                "Device complication",
                "Fall",
                "Chemo maintenance",
            ],
            "ccs_category": [
                "108",
                "100",
                "49",
                "158",
                "127",
                "96",
                "651",
                "657",
                "109",
                "237",
                "2614",
                "45",
            ],
            "ccs_description": [
                "CHF",
                "AMI",
                "Diabetes mellitus",
                "Kidney disease",
                "COPD",
                "Cardiac arrhythmia",
                "Dementia",
                "Mood disorders",
                "Stroke",
                "Device complication",
                "Struck by or against",
                "Chemo/radio",
            ],
        }
    )


def _make_ccs_icd10_pcs_df() -> pl.DataFrame:
    return pl.DataFrame(
        {
            "icd_10_pcs": ["0DTJ4ZZ", "06BK0ZZ"],
            "description": ["Resection appendix", "Excision femoral vein"],
            "ccs_category": ["74", "64"],
            "ccs_description": ["GI procedures", "Bone marrow transplant"],
        }
    )


# ═══════════════════════════════════════════════════════════════════════════════
# Module-level constants
# ═══════════════════════════════════════════════════════════════════════════════


class TestModuleConstants:
    """Verify the module-level constant definitions."""

    def test_mcc_groups_count(self) -> None:
        assert len(MCC_GROUPS) == 9

    def test_mcc_groups_values(self) -> None:
        expected = {
            "AMI",
            "ALZHEIMER",
            "AFIB",
            "CKD",
            "COPD_ASTHMA",
            "DEPRESSION",
            "DIABETES",
            "HEART_FAILURE",
            "STROKE_TIA",
        }
        assert set(MCC_GROUPS) == expected

    def test_specialty_cohorts_count(self) -> None:
        assert len(SPECIALTY_COHORTS) == 5

    def test_specialty_cohorts_contains_medicine(self) -> None:
        assert "MEDICINE" in SPECIALTY_COHORTS

    def test_specialty_cohorts_first_is_surgery_gyn(self) -> None:
        assert SPECIALTY_COHORTS[0] == "SURGERY_GYNECOLOGY"

    def test_proc_complication_ccs(self) -> None:
        assert set(_PROC_COMPLICATION_CCS) == {145, 237, 238, 257}

    def test_injury_accident_ccs_count(self) -> None:
        # MIF §3.7 lists 19 E-code categories (2601, 2602, 2604–2616, 2618–2621)
        assert len(_INJURY_ACCIDENT_CCS) == 19

    def test_injury_ccs_contains_key_categories(self) -> None:
        # Motor vehicle traffic, fire, firearm all present
        assert 2607 in _INJURY_ACCIDENT_CCS  # MVT
        assert 2604 in _INJURY_ACCIDENT_CCS  # Fire/burn
        assert 2605 in _INJURY_ACCIDENT_CCS  # Firearm


# ═══════════════════════════════════════════════════════════════════════════════
# Staging tests
# ═══════════════════════════════════════════════════════════════════════════════


class TestStgMedicalClaim:
    """Tests for stg_medical_claim staging function."""

    def test_adds_principal_diagnosis_code(self) -> None:
        medical_claim = pl.DataFrame(
            {
                "claim_id": ["CLM001", "CLM002"],
                "person_id": ["P001", "P002"],
                "claim_start_date": [date(2024, 1, 1), date(2024, 2, 1)],
                "claim_end_date": [date(2024, 1, 3), date(2024, 2, 3)],
                "hcpcs_code": pl.Series([None, None], dtype=pl.String),
                "place_of_service_code": ["21", "21"],
            }
        )
        condition = pl.DataFrame(
            {
                "claim_id": ["CLM001", "CLM001", "CLM002"],
                "normalized_code": ["I21.0", "I50.1", "E11.9"],
                "condition_rank": [1, 2, 1],
            }
        )
        result = stg_medical_claim(medical_claim, condition)
        assert "principal_diagnosis_code" in result.columns
        row1 = result.filter(pl.col("claim_id") == "CLM001")
        assert row1["principal_diagnosis_code"][0] == "I21.0"

    def test_preserves_claim_line_grain(self) -> None:
        medical_claim = pl.DataFrame(
            {
                "claim_id": ["CLM001", "CLM001"],
                "person_id": ["P001", "P001"],
                "claim_start_date": [date(2024, 1, 1), date(2024, 1, 1)],
                "claim_end_date": [date(2024, 1, 3), date(2024, 1, 3)],
                "hcpcs_code": ["99213", "71046"],
                "place_of_service_code": ["21", "21"],
            }
        )
        condition = pl.DataFrame(
            {
                "claim_id": ["CLM001"],
                "normalized_code": ["I21.0"],
                "condition_rank": [1],
            }
        )
        result = stg_medical_claim(medical_claim, condition)
        assert len(result) == 2

    def test_null_when_no_condition(self) -> None:
        medical_claim = pl.DataFrame(
            {
                "claim_id": ["CLM001"],
                "person_id": ["P001"],
                "claim_start_date": [date(2024, 1, 1)],
                "claim_end_date": [date(2024, 1, 3)],
                "hcpcs_code": pl.Series([None], dtype=pl.String),
                "place_of_service_code": ["21"],
            }
        )
        condition = pl.DataFrame(
            {
                "claim_id": pl.Series([], dtype=pl.String),
                "normalized_code": pl.Series([], dtype=pl.String),
                "condition_rank": pl.Series([], dtype=pl.Int64),
            }
        )
        result = stg_medical_claim(medical_claim, condition)
        assert len(result) == 1
        assert result["principal_diagnosis_code"][0] is None


class TestStgMedicalClaimCondition:
    """Tests for stg_medical_claim_condition staging function."""

    def test_fans_out_to_all_conditions(self) -> None:
        medical_claim = pl.DataFrame(
            {
                "claim_id": ["CLM001", "CLM001"],
                "person_id": ["P001", "P001"],
                "claim_start_date": [date(2024, 1, 1), date(2024, 1, 1)],
                "claim_end_date": [date(2024, 1, 3), date(2024, 1, 3)],
                "hcpcs_code": ["99213", "71046"],
                "place_of_service_code": ["21", "21"],
            }
        )
        condition = pl.DataFrame(
            {
                "claim_id": ["CLM001", "CLM001"],
                "normalized_code": ["I21.0", "I50.1"],
                "condition_rank": [1, 2],
            }
        )
        result = stg_medical_claim_condition(medical_claim, condition)
        assert len(result) == 2
        assert set(result["normalized_code"].to_list()) == {"I21.0", "I50.1"}

    def test_expected_columns(self) -> None:
        medical_claim = pl.DataFrame(
            {
                "claim_id": ["CLM001"],
                "person_id": ["P001"],
                "claim_start_date": [date(2024, 1, 1)],
                "claim_end_date": [date(2024, 1, 3)],
                "hcpcs_code": pl.Series([None], dtype=pl.String),
                "place_of_service_code": ["21"],
            }
        )
        condition = pl.DataFrame(
            {
                "claim_id": ["CLM001"],
                "normalized_code": ["I21.0"],
                "condition_rank": [1],
            }
        )
        result = stg_medical_claim_condition(medical_claim, condition)
        assert set(result.columns) == {
            "claim_id",
            "person_id",
            "claim_start_date",
            "normalized_code",
        }


# ═══════════════════════════════════════════════════════════════════════════════
# UAMCC tests
# ═══════════════════════════════════════════════════════════════════════════════


class TestUamccPerformancePeriod:
    """Tests for uamcc_performance_period passthrough."""

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

    def test_single_row(self, uamcc_period_df) -> None:
        result = uamcc_performance_period(uamcc_period_df)
        assert len(result) == 1

    def test_expected_columns(self, uamcc_period_df) -> None:
        result = uamcc_performance_period(uamcc_period_df)
        for col in [
            "measure_id",
            "measure_name",
            "nqf_id",
            "performance_year",
            "performance_period_begin",
            "performance_period_end",
            "lookback_period_begin",
            "lookback_period_end",
        ]:
            assert col in result.columns

    def test_measure_id_is_uamcc(self, uamcc_period_df) -> None:
        result = uamcc_performance_period(uamcc_period_df)
        assert result["measure_id"][0] == "UAMCC"

    def test_nqf_id(self, uamcc_period_df) -> None:
        result = uamcc_performance_period(uamcc_period_df)
        assert result["nqf_id"][0] == "2888"

    def test_period_dates(self, uamcc_period_df) -> None:
        result = uamcc_performance_period(uamcc_period_df)
        assert result["performance_period_begin"][0] == date(2025, 1, 1)
        assert result["performance_period_end"][0] == date(2025, 12, 31)

    def test_lookback_period(self, uamcc_period_df) -> None:
        result = uamcc_performance_period(uamcc_period_df)
        assert result["lookback_period_begin"][0] == date(2024, 1, 1)
        assert result["lookback_period_end"][0] == date(2024, 12, 31)


class TestUamccIntMccCohort:
    """Tests for uamcc_int_mcc_cohort chronic condition identification."""

    def test_returns_dataframe(
        self, stg_medical_claim_condition_df, uamcc_cohort_value_set_df
    ) -> None:
        result = uamcc_int_mcc_cohort(
            stg_medical_claim_condition_df, uamcc_cohort_value_set_df
        )
        assert isinstance(result, pl.DataFrame)

    def test_expected_columns(
        self, stg_medical_claim_condition_df, uamcc_cohort_value_set_df
    ) -> None:
        result = uamcc_int_mcc_cohort(
            stg_medical_claim_condition_df, uamcc_cohort_value_set_df
        )
        for col in [
            "person_id",
            "chronic_condition_group",
            "qualifying_code",
            "qualifying_code_date",
            "claim_count",
            "lookback_years",
        ]:
            assert col in result.columns

    def test_identifies_ami_for_p001(
        self, stg_medical_claim_condition_df, uamcc_cohort_value_set_df
    ) -> None:
        result = uamcc_int_mcc_cohort(
            stg_medical_claim_condition_df, uamcc_cohort_value_set_df
        )
        p001 = result.filter(
            (pl.col("person_id") == "P001")
            & (pl.col("chronic_condition_group") == "AMI")
        )
        assert len(p001) == 1

    def test_identifies_heart_failure_for_p001(
        self, stg_medical_claim_condition_df, uamcc_cohort_value_set_df
    ) -> None:
        result = uamcc_int_mcc_cohort(
            stg_medical_claim_condition_df, uamcc_cohort_value_set_df
        )
        p001_hf = result.filter(
            (pl.col("person_id") == "P001")
            & (pl.col("chronic_condition_group") == "HEART_FAILURE")
        )
        assert len(p001_hf) == 1

    def test_p001_has_two_conditions(
        self, stg_medical_claim_condition_df, uamcc_cohort_value_set_df
    ) -> None:
        result = uamcc_int_mcc_cohort(
            stg_medical_claim_condition_df, uamcc_cohort_value_set_df
        )
        p001 = result.filter(pl.col("person_id") == "P001")
        assert len(p001) == 2

    def test_no_match_returns_empty(self, uamcc_cohort_value_set_df) -> None:
        no_match = pl.DataFrame(
            {
                "claim_id": ["CLM999"],
                "person_id": ["P999"],
                "claim_start_date": [date(2024, 1, 1)],
                "normalized_code": ["Z99.99"],
            }
        )
        result = uamcc_int_mcc_cohort(no_match, uamcc_cohort_value_set_df)
        assert len(result) == 0

    def test_claim_count_correct(self, uamcc_cohort_value_set_df) -> None:
        # Two AMI claims for same person
        two_ami_claims = pl.DataFrame(
            {
                "claim_id": ["CLM_A", "CLM_B"],
                "person_id": ["P001", "P001"],
                "claim_start_date": [date(2024, 1, 1), date(2024, 3, 1)],
                "normalized_code": ["I21.0", "I21.1"],
            }
        )
        result = uamcc_int_mcc_cohort(two_ami_claims, uamcc_cohort_value_set_df)
        ami_row = result.filter(pl.col("chronic_condition_group") == "AMI")
        assert ami_row["claim_count"][0] == 2


class TestUamccIntDenominator:
    """Tests for uamcc_int_denominator eligibility filtering."""

    @pytest.fixture
    def mcc_cohort_df(self) -> pl.DataFrame:
        """Cohort with two or more conditions for each person."""
        return pl.DataFrame(
            {
                "person_id": [
                    "P001",
                    "P001",
                    "P002",
                    "P002",
                    "P003",
                    "P003",  # age 60 — should be excluded by age filter
                    "P004",
                    "P004",  # deceased mid-year, still included in denominator
                    "P005",
                    "P005",
                ],
                "chronic_condition_group": [
                    "AMI",
                    "HEART_FAILURE",
                    "DIABETES",
                    "CKD",
                    "COPD_ASTHMA",
                    "AFIB",
                    "COPD_ASTHMA",
                    "AFIB",
                    "ALZHEIMER",
                    "DEPRESSION",
                ],
                "qualifying_code": [
                    "I21.0",
                    "I50.1",
                    "E11.9",
                    "N18.3",
                    "J44.1",
                    "I48.0",
                    "J44.1",
                    "I48.0",
                    "G30.0",
                    "F32.0",
                ],
                "qualifying_code_date": [date(2024, 3, 1)] * 10,
                "claim_count": [1] * 10,
                "lookback_years": [1] * 10,
            }
        )

    def test_returns_dataframe(
        self, mcc_cohort_df, patient_df, uamcc_period_df
    ) -> None:
        result = uamcc_int_denominator(mcc_cohort_df, patient_df, uamcc_period_df)
        assert isinstance(result, pl.DataFrame)

    def test_expected_columns(self, mcc_cohort_df, patient_df, uamcc_period_df) -> None:
        result = uamcc_int_denominator(mcc_cohort_df, patient_df, uamcc_period_df)
        for col in ["person_id", "age_at_period_start", "chronic_condition_count"]:
            assert col in result.columns

    def test_excludes_patient_under_66(
        self, mcc_cohort_df, patient_df, uamcc_period_df
    ) -> None:
        result = uamcc_int_denominator(mcc_cohort_df, patient_df, uamcc_period_df)
        # P003 is age 60 — must be excluded
        assert "P003" not in result["person_id"].to_list()

    def test_includes_patient_over_66(
        self, mcc_cohort_df, patient_df, uamcc_period_df
    ) -> None:
        result = uamcc_int_denominator(mcc_cohort_df, patient_df, uamcc_period_df)
        # P001 (age ~79), P002, P004, P005 all ≥66
        included = result["person_id"].to_list()
        assert "P001" in included
        assert "P002" in included

    def test_requires_two_or_more_conditions(self, patient_df, uamcc_period_df) -> None:
        one_condition = pl.DataFrame(
            {
                "person_id": ["P001"],
                "chronic_condition_group": ["AMI"],
                "qualifying_code": ["I21.0"],
                "qualifying_code_date": [date(2024, 3, 1)],
                "claim_count": [1],
                "lookback_years": [1],
            }
        )
        result = uamcc_int_denominator(one_condition, patient_df, uamcc_period_df)
        assert len(result) == 0

    def test_chronic_condition_count_correct(
        self, mcc_cohort_df, patient_df, uamcc_period_df
    ) -> None:
        result = uamcc_int_denominator(mcc_cohort_df, patient_df, uamcc_period_df)
        p001 = result.filter(pl.col("person_id") == "P001")
        assert p001["chronic_condition_count"][0] == 2

    def test_empty_cohort_returns_empty(self, patient_df, uamcc_period_df) -> None:
        empty = pl.DataFrame(
            {
                "person_id": pl.Series([], dtype=pl.String),
                "chronic_condition_group": pl.Series([], dtype=pl.String),
                "qualifying_code": pl.Series([], dtype=pl.String),
                "qualifying_code_date": pl.Series([], dtype=pl.Date),
                "claim_count": pl.Series([], dtype=pl.Int32),
                "lookback_years": pl.Series([], dtype=pl.Int32),
            }
        )
        result = uamcc_int_denominator(empty, patient_df, uamcc_period_df)
        assert len(result) == 0


class TestUamccIntDenominatorExclusion:
    """Tests for uamcc_int_denominator_exclusion flag logic."""

    @pytest.fixture
    def denominator_df(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "person_id": ["P001", "P002", "P004"],
                "age_at_period_start": [79, 76, 74],
                "chronic_condition_count": [2, 2, 2],
            }
        )

    def test_returns_dataframe(self, denominator_df, patient_df) -> None:
        result = uamcc_int_denominator_exclusion(denominator_df, patient_df)
        assert isinstance(result, pl.DataFrame)

    def test_expected_columns(self, denominator_df, patient_df) -> None:
        result = uamcc_int_denominator_exclusion(denominator_df, patient_df)
        for col in [
            "person_id",
            "exclusion_reason",
            "voluntary_alignment_after_period_start",
            "missing_prior_year_enrollment",
            "missing_measurement_year_enrollment",
            "in_hospice",
            "no_aco_visit",
            "no_time_at_risk",
        ]:
            assert col in result.columns

    def test_deceased_at_period_start_flagged(self, denominator_df, patient_df) -> None:
        # P004 has death_date set in patient_df → no_time_at_risk
        result = uamcc_int_denominator_exclusion(denominator_df, patient_df)
        p004 = result.filter(pl.col("person_id") == "P004")
        if len(p004) > 0:
            assert p004["no_time_at_risk"][0] == 1

    def test_living_beneficiary_not_excluded_by_default(
        self, denominator_df, patient_df
    ) -> None:
        # P001 and P002 are alive — no_time_at_risk should be 0
        result = uamcc_int_denominator_exclusion(denominator_df, patient_df)
        # They should either not appear (no exclusion triggered) or appear with 0
        alive = result.filter(pl.col("person_id").is_in(["P001", "P002"]))
        for row in alive.iter_rows(named=True):
            assert row["no_time_at_risk"] == 0


class TestUamccIntPlannedAdmission:
    """Tests for uamcc_int_planned_admission PAA algorithm."""

    def test_returns_dataframe(self, stg_medical_claim_df) -> None:
        result = uamcc_int_planned_admission(
            stg_medical_claim_df,
            _make_paa1_df(),
            _make_paa2_df(),
            _make_paa3_df(),
            _make_paa4_df(),
            _make_ccs_icd10_cm_df(),
            _make_ccs_icd10_pcs_df(),
        )
        assert isinstance(result, pl.DataFrame)

    def test_expected_columns(self, stg_medical_claim_df) -> None:
        result = uamcc_int_planned_admission(
            stg_medical_claim_df,
            _make_paa1_df(),
            _make_paa2_df(),
            _make_paa3_df(),
            _make_paa4_df(),
            _make_ccs_icd10_cm_df(),
            _make_ccs_icd10_pcs_df(),
        )
        for col in [
            "claim_id",
            "person_id",
            "admission_date",
            "is_planned",
            "planned_rule",
        ]:
            assert col in result.columns

    def test_rule2_planned_diagnosis(self) -> None:
        """Claim with principal diagnosis CCS 45 (chemo) → Rule 2 planned."""
        chemo_claim = pl.DataFrame(
            {
                "claim_id": ["CLM_CHEMO"],
                "person_id": ["P001"],
                "claim_start_date": [date(2025, 3, 1)],
                "claim_end_date": [date(2025, 3, 3)],
                "principal_diagnosis_code": ["Z51.11"],
                "hcpcs_code": pl.Series([None], dtype=pl.String),
                "place_of_service_code": ["21"],
            }
        )
        result = uamcc_int_planned_admission(
            chemo_claim,
            _make_paa1_df(),
            _make_paa2_df(),
            _make_paa3_df(),
            _make_paa4_df(),
            _make_ccs_icd10_cm_df(),
            _make_ccs_icd10_pcs_df(),
        )
        row = result.filter(pl.col("claim_id") == "CLM_CHEMO")
        assert row["is_planned"][0] == 1
        assert row["planned_rule"][0] == "RULE2"

    def test_unplanned_acute_mi(self, stg_medical_claim_df) -> None:
        """AMI claim (I21.0) is an acute diagnosis → unplanned."""
        result = uamcc_int_planned_admission(
            stg_medical_claim_df,
            _make_paa1_df(),
            _make_paa2_df(),
            _make_paa3_df(),
            _make_paa4_df(),
            _make_ccs_icd10_cm_df(),
            _make_ccs_icd10_pcs_df(),
        )
        ami_claim = result.filter(pl.col("claim_id") == "CLM001")
        assert ami_claim["is_planned"][0] == 0
        assert ami_claim["planned_rule"][0] is None

    def test_unplanned_row_count_matches_input(self, stg_medical_claim_df) -> None:
        result = uamcc_int_planned_admission(
            stg_medical_claim_df,
            _make_paa1_df(),
            _make_paa2_df(),
            _make_paa3_df(),
            _make_paa4_df(),
            _make_ccs_icd10_cm_df(),
            _make_ccs_icd10_pcs_df(),
        )
        assert len(result) == len(stg_medical_claim_df)


class TestUamccIntOutcomeExclusion:
    """Tests for uamcc_int_outcome_exclusion CCS-based exclusion flagging."""

    @pytest.fixture
    def device_complication_claim(self) -> pl.DataFrame:
        """Claim with CCS 237 (device complication) as principal diagnosis."""
        return pl.DataFrame(
            {
                "claim_id": ["CLM_COMPL"],
                "person_id": ["P001"],
                "claim_start_date": [date(2025, 4, 1)],
                "claim_end_date": [date(2025, 4, 3)],
                "principal_diagnosis_code": ["T82.7XXA"],
                "hcpcs_code": pl.Series([None], dtype=pl.String),
                "place_of_service_code": ["21"],
            }
        )

    @pytest.fixture
    def injury_claim(self) -> pl.DataFrame:
        """Claim with CCS 2614 (struck by) as principal diagnosis."""
        return pl.DataFrame(
            {
                "claim_id": ["CLM_INJURY"],
                "person_id": ["P001"],
                "claim_start_date": [date(2025, 5, 1)],
                "claim_end_date": [date(2025, 5, 2)],
                "principal_diagnosis_code": ["W19.XXXA"],
                "hcpcs_code": pl.Series([None], dtype=pl.String),
                "place_of_service_code": ["21"],
            }
        )

    @pytest.fixture
    def exclusions_value_set_df(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "exclusion_category": [
                    "Complications of procedures or surgeries",
                    "Complications of procedures or surgeries",
                    "Accidents/Injuries",
                    "Accidents/Injuries",
                ],
                "code_type": ["CCS", "CCS", "CCS", "CCS"],
                "category_or_code": ["237", "238", "2614", "2607"],
                "description": [
                    "Device complication",
                    "Surgical complication",
                    "Struck by",
                    "MVT",
                ],
                "procedure_or_diagnosis": ["Diagnosis"] * 4,
            }
        )

    @pytest.fixture
    def empty_planned_df(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "claim_id": pl.Series([], dtype=pl.String),
                "is_planned": pl.Series([], dtype=pl.Int32),
            }
        )

    def test_flags_procedure_complication(
        self, device_complication_claim, exclusions_value_set_df, empty_planned_df
    ) -> None:
        result = uamcc_int_outcome_exclusion(
            device_complication_claim,
            empty_planned_df,
            exclusions_value_set_df,
            _make_ccs_icd10_cm_df(),
        )
        assert len(result) == 1
        assert result["is_procedure_complication"][0] == 1

    def test_flags_injury(
        self, injury_claim, exclusions_value_set_df, empty_planned_df
    ) -> None:
        result = uamcc_int_outcome_exclusion(
            injury_claim,
            empty_planned_df,
            exclusions_value_set_df,
            _make_ccs_icd10_cm_df(),
        )
        assert len(result) == 1
        assert result["is_injury_or_accident"][0] == 1

    def test_non_excluded_claim_not_returned(
        self, stg_medical_claim_df, exclusions_value_set_df, empty_planned_df
    ) -> None:
        # Heart failure (CCS 108) — not in exclusion list
        hf_only = stg_medical_claim_df.filter(pl.col("claim_id") == "CLM002")
        result = uamcc_int_outcome_exclusion(
            hf_only,
            empty_planned_df,
            exclusions_value_set_df,
            _make_ccs_icd10_cm_df(),
        )
        assert len(result) == 0

    def test_planned_claim_flagged(
        self, stg_medical_claim_df, exclusions_value_set_df
    ) -> None:
        planned_df = pl.DataFrame(
            {
                "claim_id": ["CLM001"],
                "is_planned": [1],
            }
        )
        result = uamcc_int_outcome_exclusion(
            stg_medical_claim_df.filter(pl.col("claim_id") == "CLM001"),
            planned_df,
            exclusions_value_set_df,
            _make_ccs_icd10_cm_df(),
        )
        assert len(result) == 1
        assert result["is_planned"][0] == 1

    def test_expected_columns(
        self, device_complication_claim, exclusions_value_set_df, empty_planned_df
    ) -> None:
        result = uamcc_int_outcome_exclusion(
            device_complication_claim,
            empty_planned_df,
            exclusions_value_set_df,
            _make_ccs_icd10_cm_df(),
        )
        for col in [
            "claim_id",
            "person_id",
            "admission_date",
            "is_planned",
            "from_snf_or_rehab",
            "in_buffer_period",
            "in_hospice",
            "is_procedure_complication",
            "is_injury_or_accident",
            "before_first_aco_visit",
            "ccs_diagnosis_category",
        ]:
            assert col in result.columns


class TestUamccIntPersonTime:
    """Tests for uamcc_int_person_time at-risk calculation."""

    @pytest.fixture
    def denominator_df(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "person_id": ["P001", "P002"],
                "age_at_period_start": [79, 76],
                "chronic_condition_count": [2, 2],
            }
        )

    def test_returns_dataframe(
        self, denominator_df, encounter_df, uamcc_period_df
    ) -> None:
        result = uamcc_int_person_time(denominator_df, encounter_df, uamcc_period_df)
        assert isinstance(result, pl.DataFrame)

    def test_expected_columns(
        self, denominator_df, encounter_df, uamcc_period_df
    ) -> None:
        result = uamcc_int_person_time(denominator_df, encounter_df, uamcc_period_df)
        for col in [
            "person_id",
            "at_risk_days",
            "person_years",
            "days_in_hospital",
            "days_in_snf_rehab",
            "days_in_buffer",
            "days_in_hospice",
        ]:
            assert col in result.columns

    def test_at_risk_days_positive(
        self, denominator_df, encounter_df, uamcc_period_df
    ) -> None:
        result = uamcc_int_person_time(denominator_df, encounter_df, uamcc_period_df)
        assert (result["at_risk_days"] >= 0).all()

    def test_person_years_derived_from_days(
        self, denominator_df, encounter_df, uamcc_period_df
    ) -> None:
        result = uamcc_int_person_time(denominator_df, encounter_df, uamcc_period_df)
        for row in result.iter_rows(named=True):
            expected = row["at_risk_days"] / 365.25
            assert abs(row["person_years"] - expected) < 0.01

    def test_hospital_days_reduce_at_risk(
        self, denominator_df, encounter_df, uamcc_period_df
    ) -> None:
        result = uamcc_int_person_time(denominator_df, encounter_df, uamcc_period_df)
        # P001 has 2 acute inpatient stays: 4 days + 3 days = 7 days in hospital
        p001 = result.filter(pl.col("person_id") == "P001")
        assert p001["days_in_hospital"][0] == 7

    def test_one_row_per_person(
        self, denominator_df, encounter_df, uamcc_period_df
    ) -> None:
        result = uamcc_int_person_time(denominator_df, encounter_df, uamcc_period_df)
        assert result["person_id"].n_unique() == len(denominator_df)


# ═══════════════════════════════════════════════════════════════════════════════
# ACR tests
# ═══════════════════════════════════════════════════════════════════════════════


class TestAcrPerformancePeriod:
    """Tests for acr_performance_period passthrough."""

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

    def test_single_row(self, acr_period_df) -> None:
        assert len(acr_performance_period(acr_period_df)) == 1

    def test_measure_id_is_acr(self, acr_period_df) -> None:
        result = acr_performance_period(acr_period_df)
        assert result["measure_id"][0] == "ACR"

    def test_nqf_id(self, acr_period_df) -> None:
        result = acr_performance_period(acr_period_df)
        assert result["nqf_id"][0] == "1789"

    def test_expected_columns(self, acr_period_df) -> None:
        result = acr_performance_period(acr_period_df)
        for col in [
            "measure_id",
            "measure_name",
            "nqf_id",
            "performance_year",
            "performance_period_begin",
            "performance_period_end",
        ]:
            assert col in result.columns


class TestAcrIntIndexAdmission:
    """Tests for acr_int_index_admission denominator building."""

    @pytest.fixture
    def excl_vs_df(self) -> pl.DataFrame:
        """ACR exclusion value set — psychiatric stays excluded."""
        return pl.DataFrame(
            {
                "ccs_diagnosis_category": ["650", "651"],
                "description": ["Adjustment disorders", "Anxiety disorders"],
                "exclusion_category": ["Psychiatric"] * 2,
            }
        )

    def test_returns_dataframe(self, encounter_df, excl_vs_df) -> None:
        result = acr_int_index_admission(
            encounter_df, excl_vs_df, _make_ccs_icd10_cm_df()
        )
        assert isinstance(result, pl.DataFrame)

    def test_filters_to_acute_inpatient(self, encounter_df, excl_vs_df) -> None:
        result = acr_int_index_admission(
            encounter_df, excl_vs_df, _make_ccs_icd10_cm_df()
        )
        # ENC005 is outpatient — should not appear
        assert "ENC005" not in result["encounter_id"].to_list()

    def test_expected_columns(self, encounter_df, excl_vs_df) -> None:
        result = acr_int_index_admission(
            encounter_df, excl_vs_df, _make_ccs_icd10_cm_df()
        )
        for col in [
            "encounter_id",
            "person_id",
            "admission_date",
            "discharge_date",
            "principal_diagnosis_code",
            "ccs_diagnosis_category",
            "exclusion_flag",
            "exclusion_reason",
        ]:
            assert col in result.columns

    def test_excluded_diagnosis_flagged(self) -> None:
        """Psychiatric diagnosis CCS 651 → exclusion_flag = 1."""
        psych_encounter = pl.DataFrame(
            {
                "encounter_id": ["ENC_PSY"],
                "person_id": ["P001"],
                "encounter_type": ["acute inpatient"],
                "encounter_start_date": [date(2025, 2, 1)],
                "encounter_end_date": [date(2025, 2, 5)],
                "length_of_stay": [4],
                "discharge_disposition_code": ["01"],
                "facility_id": ["H001"],
                "primary_diagnosis_code": ["G30.0"],  # CCS 651 in our test map
                "ccs_diagnosis_category": ["651"],
                "drg_code_type": ["MS-DRG"],
                "drg_code": ["897"],
                "encounter_group": ["claims"],
            },
            schema={
                "encounter_id": pl.String,
                "person_id": pl.String,
                "encounter_type": pl.String,
                "encounter_start_date": pl.Date,
                "encounter_end_date": pl.Date,
                "length_of_stay": pl.Int32,
                "discharge_disposition_code": pl.String,
                "facility_id": pl.String,
                "primary_diagnosis_code": pl.String,
                "ccs_diagnosis_category": pl.String,
                "drg_code_type": pl.String,
                "drg_code": pl.String,
                "encounter_group": pl.String,
            },
        )
        excl_vs = pl.DataFrame(
            {
                "ccs_diagnosis_category": ["651"],
                "description": ["Anxiety disorders"],
                "exclusion_category": ["Psychiatric"],
            }
        )
        result = acr_int_index_admission(
            psych_encounter, excl_vs, _make_ccs_icd10_cm_df()
        )
        assert result["exclusion_flag"][0] == 1

    def test_non_excluded_encounter_has_zero_flag(
        self, encounter_df, excl_vs_df
    ) -> None:
        result = acr_int_index_admission(
            encounter_df, excl_vs_df, _make_ccs_icd10_cm_df()
        )
        non_excluded = result.filter(pl.col("exclusion_flag") == 0)
        assert len(non_excluded) > 0


class TestAcrIntSpecialtyCohort:
    """Tests for acr_int_specialty_cohort CCS-based assignment."""

    @pytest.fixture
    def cohort_ccs_df(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "ccs_category": ["108", "100", "127"],
                "description": ["CHF", "AMI", "COPD"],
                "specialty_cohort": [
                    "CARDIORESPIRATORY",
                    "CARDIOVASCULAR",
                    "CARDIORESPIRATORY",
                ],
                "procedure_or_diagnosis": ["Diagnosis", "Diagnosis", "Diagnosis"],
                "principal_diagnosis_code": [1, 1, 1],
                "procedure_code": [0, 0, 0],
            }
        )

    @pytest.fixture
    def cohort_icd10_df(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "icd_10_pcs": ["0FT40ZZ"],
                "description": ["Resection gallbladder"],
                "associated_ccs_category": ["79"],
                "specialty_cohort": ["SURGERY_GYNECOLOGY"],
            }
        )

    @pytest.fixture
    def procedure_df(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "encounter_id": ["ENC_SURG"],
                "person_id": ["P001"],
                "normalized_code": ["0FT40ZZ"],
                "procedure_date": [date(2025, 3, 1)],
                "source_code_type": ["ICD-10-PCS"],
                "source_code": ["0FT40ZZ"],
            }
        )

    @pytest.fixture
    def index_admission_df(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "encounter_id": ["ENC001", "ENC002", "ENC003", "ENC_SURG"],
                "person_id": ["P001", "P001", "P002", "P001"],
                "admission_date": [
                    date(2025, 2, 1),
                    date(2025, 4, 5),
                    date(2025, 3, 1),
                    date(2025, 6, 1),
                ],
                "discharge_date": [
                    date(2025, 2, 5),
                    date(2025, 4, 8),
                    date(2025, 3, 4),
                    date(2025, 6, 4),
                ],
                "principal_diagnosis_code": ["I50.1", "I21.0", "E11.9", "K80.00"],
                "ccs_diagnosis_category": ["108", "100", "49", "149"],
                "exclusion_flag": [0, 0, 0, 0],
                "exclusion_reason": [None, None, None, None],
                "discharge_disposition_code": ["01"] * 4,
                "facility_id": ["H001"] * 4,
                "drg_code_type": ["MS-DRG"] * 4,
                "drg_code": ["291", "282", "637", "395"],
            },
            schema={
                "encounter_id": pl.String,
                "person_id": pl.String,
                "admission_date": pl.Date,
                "discharge_date": pl.Date,
                "principal_diagnosis_code": pl.String,
                "ccs_diagnosis_category": pl.String,
                "exclusion_flag": pl.Int32,
                "exclusion_reason": pl.String,
                "discharge_disposition_code": pl.String,
                "facility_id": pl.String,
                "drg_code_type": pl.String,
                "drg_code": pl.String,
            },
        )

    def test_returns_dataframe(
        self, index_admission_df, cohort_ccs_df, cohort_icd10_df, procedure_df
    ) -> None:
        result = acr_int_specialty_cohort(
            index_admission_df, cohort_ccs_df, cohort_icd10_df, procedure_df
        )
        assert isinstance(result, pl.DataFrame)

    def test_expected_columns(
        self, index_admission_df, cohort_ccs_df, cohort_icd10_df, procedure_df
    ) -> None:
        result = acr_int_specialty_cohort(
            index_admission_df, cohort_ccs_df, cohort_icd10_df, procedure_df
        )
        for col in ["encounter_id", "specialty_cohort", "cohort_assignment_rule"]:
            assert col in result.columns

    def test_cardiorespiratory_from_ccs(
        self, index_admission_df, cohort_ccs_df, cohort_icd10_df, procedure_df
    ) -> None:
        result = acr_int_specialty_cohort(
            index_admission_df, cohort_ccs_df, cohort_icd10_df, procedure_df
        )
        enc1 = result.filter(pl.col("encounter_id") == "ENC001")
        assert enc1["specialty_cohort"][0] == "CARDIORESPIRATORY"

    def test_surgery_gyn_from_icd_pcs(
        self, index_admission_df, cohort_ccs_df, cohort_icd10_df, procedure_df
    ) -> None:
        result = acr_int_specialty_cohort(
            index_admission_df, cohort_ccs_df, cohort_icd10_df, procedure_df
        )
        enc_surg = result.filter(pl.col("encounter_id") == "ENC_SURG")
        assert enc_surg["specialty_cohort"][0] == "SURGERY_GYNECOLOGY"
        assert enc_surg["cohort_assignment_rule"][0] == "ICD10_PCS"

    def test_default_medicine(
        self, index_admission_df, cohort_ccs_df, cohort_icd10_df
    ) -> None:
        """Encounter with CCS not in cohort map → MEDICINE default."""
        empty_proc = pl.DataFrame(
            {
                "encounter_id": pl.Series([], dtype=pl.String),
                "person_id": pl.Series([], dtype=pl.String),
                "normalized_code": pl.Series([], dtype=pl.String),
                "procedure_date": pl.Series([], dtype=pl.Date),
                "source_code_type": pl.Series([], dtype=pl.String),
                "source_code": pl.Series([], dtype=pl.String),
            }
        )
        result = acr_int_specialty_cohort(
            index_admission_df.filter(pl.col("encounter_id") == "ENC003"),
            cohort_ccs_df,
            cohort_icd10_df,
            empty_proc,
        )
        # CCS 49 (Diabetes) is not in our small cohort_ccs_df → MEDICINE
        assert result["specialty_cohort"][0] == "MEDICINE"

    def test_surgery_wins_over_diagnosis_cohort(
        self, index_admission_df, cohort_ccs_df, cohort_icd10_df, procedure_df
    ) -> None:
        """ENC_SURG: has SURGERY_GYNECOLOGY procedure; diagnosis CCS not mapped → still SURGERY_GYNECOLOGY."""
        result = acr_int_specialty_cohort(
            index_admission_df, cohort_ccs_df, cohort_icd10_df, procedure_df
        )
        enc_surg = result.filter(pl.col("encounter_id") == "ENC_SURG")
        assert enc_surg["specialty_cohort"][0] == "SURGERY_GYNECOLOGY"


class TestAcrIntPlannedReadmission:
    """Tests for acr_int_planned_readmission PAA classification."""

    @pytest.fixture
    def index_admission_df(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "encounter_id": ["ENC001"],
                "person_id": ["P001"],
                "admission_date": [date(2025, 2, 1)],
                "discharge_date": [date(2025, 2, 5)],
                "exclusion_flag": [0],
                "exclusion_reason": [None],
                "principal_diagnosis_code": ["I50.1"],
                "ccs_diagnosis_category": ["108"],
                "discharge_disposition_code": ["01"],
                "facility_id": ["H001"],
                "drg_code_type": ["MS-DRG"],
                "drg_code": ["291"],
            },
            schema={
                "encounter_id": pl.String,
                "person_id": pl.String,
                "admission_date": pl.Date,
                "discharge_date": pl.Date,
                "exclusion_flag": pl.Int32,
                "exclusion_reason": pl.String,
                "principal_diagnosis_code": pl.String,
                "ccs_diagnosis_category": pl.String,
                "discharge_disposition_code": pl.String,
                "facility_id": pl.String,
                "drg_code_type": pl.String,
                "drg_code": pl.String,
            },
        )

    @pytest.fixture
    def readmission_encounter_df(self) -> pl.DataFrame:
        """Two candidate readmissions: one within 30 days, one beyond."""
        return pl.DataFrame(
            {
                "encounter_id": ["ENC_RA1", "ENC_RA2", "ENC_RA3"],
                "person_id": ["P001", "P001", "P001"],
                "encounter_type": ["acute inpatient"] * 3,
                "encounter_start_date": [
                    date(2025, 2, 20),  # 15 days after discharge → within 30
                    date(2025, 3, 20),  # 43 days → beyond 30
                    date(2025, 2, 28),  # chemo maintenance → planned
                ],
                "encounter_end_date": [
                    date(2025, 2, 23),
                    date(2025, 3, 23),
                    date(2025, 3, 2),
                ],
                "length_of_stay": [3, 3, 2],
                "discharge_disposition_code": ["01", "01", "01"],
                "facility_id": ["H001"] * 3,
                "primary_diagnosis_code": ["I21.0", "I21.0", "Z51.11"],
                "ccs_diagnosis_category": ["100", "100", "45"],
                "drg_code_type": ["MS-DRG"] * 3,
                "drg_code": ["282", "282", "834"],
                "encounter_group": ["claims"] * 3,
            },
            schema={
                "encounter_id": pl.String,
                "person_id": pl.String,
                "encounter_type": pl.String,
                "encounter_start_date": pl.Date,
                "encounter_end_date": pl.Date,
                "length_of_stay": pl.Int32,
                "discharge_disposition_code": pl.String,
                "facility_id": pl.String,
                "primary_diagnosis_code": pl.String,
                "ccs_diagnosis_category": pl.String,
                "drg_code_type": pl.String,
                "drg_code": pl.String,
                "encounter_group": pl.String,
            },
        )

    def test_returns_dataframe(
        self, index_admission_df, readmission_encounter_df
    ) -> None:
        result = acr_int_planned_readmission(
            readmission_encounter_df,
            index_admission_df,
            _make_paa1_df(),
            _make_paa2_df(),
            _make_paa3_df(),
            _make_paa4_df(),
            _make_ccs_icd10_cm_df(),
            _make_ccs_icd10_pcs_df(),
        )
        assert isinstance(result, pl.DataFrame)

    def test_within_30_days_flagged(
        self, index_admission_df, readmission_encounter_df
    ) -> None:
        result = acr_int_planned_readmission(
            readmission_encounter_df,
            index_admission_df,
            _make_paa1_df(),
            _make_paa2_df(),
            _make_paa3_df(),
            _make_paa4_df(),
            _make_ccs_icd10_cm_df(),
            _make_ccs_icd10_pcs_df(),
        )
        # ENC_RA1 is 15 days after discharge → should appear with is_within_30_days=1
        ra1 = result.filter(pl.col("readmission_encounter_id") == "ENC_RA1")
        assert len(ra1) == 1
        assert ra1["is_within_30_days"][0] == 1

    def test_beyond_30_days_excluded(
        self, index_admission_df, readmission_encounter_df
    ) -> None:
        result = acr_int_planned_readmission(
            readmission_encounter_df,
            index_admission_df,
            _make_paa1_df(),
            _make_paa2_df(),
            _make_paa3_df(),
            _make_paa4_df(),
            _make_ccs_icd10_cm_df(),
            _make_ccs_icd10_pcs_df(),
        )
        # ENC_RA2 is 43 days after discharge → should not appear
        ra2 = result.filter(pl.col("readmission_encounter_id") == "ENC_RA2")
        assert len(ra2) == 0

    def test_planned_readmission_not_in_numerator(
        self, index_admission_df, readmission_encounter_df
    ) -> None:
        result = acr_int_planned_readmission(
            readmission_encounter_df,
            index_admission_df,
            _make_paa1_df(),
            _make_paa2_df(),
            _make_paa3_df(),
            _make_paa4_df(),
            _make_ccs_icd10_cm_df(),
            _make_ccs_icd10_pcs_df(),
        )
        # ENC_RA3 has chemo diagnosis (CCS 45) → PAA Rule 2 → planned
        ra3 = result.filter(pl.col("readmission_encounter_id") == "ENC_RA3")
        if len(ra3) > 0:
            assert ra3["unplanned_readmission_flag"][0] == 0

    def test_unplanned_readmission_in_numerator(
        self, index_admission_df, readmission_encounter_df
    ) -> None:
        result = acr_int_planned_readmission(
            readmission_encounter_df,
            index_admission_df,
            _make_paa1_df(),
            _make_paa2_df(),
            _make_paa3_df(),
            _make_paa4_df(),
            _make_ccs_icd10_cm_df(),
            _make_ccs_icd10_pcs_df(),
        )
        ra1 = result.filter(pl.col("readmission_encounter_id") == "ENC_RA1")
        assert ra1["unplanned_readmission_flag"][0] == 1

    def test_expected_columns(
        self, index_admission_df, readmission_encounter_df
    ) -> None:
        result = acr_int_planned_readmission(
            readmission_encounter_df,
            index_admission_df,
            _make_paa1_df(),
            _make_paa2_df(),
            _make_paa3_df(),
            _make_paa4_df(),
            _make_ccs_icd10_cm_df(),
            _make_ccs_icd10_pcs_df(),
        )
        for col in [
            "index_encounter_id",
            "readmission_encounter_id",
            "person_id",
            "index_discharge_date",
            "readmission_date",
            "days_to_readmission",
            "is_within_30_days",
            "is_planned",
            "planned_rule",
            "is_psychiatric_or_rehab",
            "unplanned_readmission_flag",
        ]:
            assert col in result.columns


# ═══════════════════════════════════════════════════════════════════════════════
# HWR tests
# ═══════════════════════════════════════════════════════════════════════════════


class TestHwrPerformancePeriod:
    """Tests for hwr_performance_period passthrough."""

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

    def test_single_row(self, hwr_period_df) -> None:
        assert len(hwr_performance_period(hwr_period_df)) == 1

    def test_measure_id_is_hwr(self, hwr_period_df) -> None:
        result = hwr_performance_period(hwr_period_df)
        assert result["measure_id"][0] == "HWR"

    def test_expected_columns(self, hwr_period_df) -> None:
        result = hwr_performance_period(hwr_period_df)
        for col in [
            "measure_id",
            "measure_name",
            "performance_year",
            "performance_period_begin",
            "performance_period_end",
        ]:
            assert col in result.columns

    def test_performance_year(self, hwr_period_df) -> None:
        result = hwr_performance_period(hwr_period_df)
        assert result["performance_year"][0] == 2025


class TestHwrIntDenominator:
    """Tests for hwr_int_denominator — MIPS HWR denominator building."""

    @pytest.fixture
    def hwr_specialty_cohort_df(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "ccs_category": ["108", "100"],
                "description": ["CHF", "AMI"],
                "specialty_cohort": ["CARDIORESPIRATORY", "CARDIOVASCULAR"],
                "procedure_or_diagnosis": ["Diagnosis", "Diagnosis"],
                "principal_diagnosis_code": [1, 1],
                "procedure_code": [0, 0],
            }
        )

    @pytest.fixture
    def hwr_surg_gyn_df(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "icd_10_pcs": ["0FT40ZZ"],
                "description": ["Resection gallbladder"],
                "associated_ccs_category": ["79"],
                "specialty_cohort": ["SURGERY_GYNECOLOGY"],
            }
        )

    @pytest.fixture
    def hwr_excl_df(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "ccs_diagnosis_category": ["650"],
                "description": ["Adjustment disorders"],
                "exclusion_category": ["Psychiatric"],
            }
        )

    @pytest.fixture
    def empty_procedure_df(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "encounter_id": pl.Series([], dtype=pl.String),
                "person_id": pl.Series([], dtype=pl.String),
                "normalized_code": pl.Series([], dtype=pl.String),
                "procedure_date": pl.Series([], dtype=pl.Date),
                "source_code_type": pl.Series([], dtype=pl.String),
                "source_code": pl.Series([], dtype=pl.String),
            }
        )

    def test_returns_dataframe(
        self,
        encounter_df,
        hwr_excl_df,
        hwr_specialty_cohort_df,
        hwr_surg_gyn_df,
        empty_procedure_df,
    ) -> None:
        result = hwr_int_denominator(
            encounter_df,
            hwr_excl_df,
            hwr_specialty_cohort_df,
            hwr_surg_gyn_df,
            empty_procedure_df,
        )
        assert isinstance(result, pl.DataFrame)

    def test_filters_to_acute_inpatient(
        self,
        encounter_df,
        hwr_excl_df,
        hwr_specialty_cohort_df,
        hwr_surg_gyn_df,
        empty_procedure_df,
    ) -> None:
        result = hwr_int_denominator(
            encounter_df,
            hwr_excl_df,
            hwr_specialty_cohort_df,
            hwr_surg_gyn_df,
            empty_procedure_df,
        )
        # ENC005 (outpatient) must not appear
        assert "ENC005" not in result["encounter_id"].to_list()

    def test_expected_columns(
        self,
        encounter_df,
        hwr_excl_df,
        hwr_specialty_cohort_df,
        hwr_surg_gyn_df,
        empty_procedure_df,
    ) -> None:
        result = hwr_int_denominator(
            encounter_df,
            hwr_excl_df,
            hwr_specialty_cohort_df,
            hwr_surg_gyn_df,
            empty_procedure_df,
        )
        for col in [
            "encounter_id",
            "person_id",
            "admission_date",
            "discharge_date",
            "specialty_cohort",
            "exclusion_flag",
            "exclusion_reason",
            "attributed_tin",
            "attribution_role",
        ]:
            assert col in result.columns

    def test_cardiorespiratory_cohort_assigned(
        self,
        encounter_df,
        hwr_excl_df,
        hwr_specialty_cohort_df,
        hwr_surg_gyn_df,
        empty_procedure_df,
    ) -> None:
        result = hwr_int_denominator(
            encounter_df,
            hwr_excl_df,
            hwr_specialty_cohort_df,
            hwr_surg_gyn_df,
            empty_procedure_df,
        )
        enc1 = result.filter(pl.col("encounter_id") == "ENC001")
        # ENC001 has CCS 108 (CHF) → CARDIORESPIRATORY
        assert enc1["specialty_cohort"][0] == "CARDIORESPIRATORY"

    def test_default_medicine_when_no_ccs_match(
        self,
        encounter_df,
        hwr_excl_df,
        hwr_specialty_cohort_df,
        hwr_surg_gyn_df,
        empty_procedure_df,
    ) -> None:
        result = hwr_int_denominator(
            encounter_df,
            hwr_excl_df,
            hwr_specialty_cohort_df,
            hwr_surg_gyn_df,
            empty_procedure_df,
        )
        # ENC003 has CCS 49 (Diabetes) → not in small specialty cohort → MEDICINE
        enc3 = result.filter(pl.col("encounter_id") == "ENC003")
        assert enc3["specialty_cohort"][0] == "MEDICINE"

    def test_non_excluded_encounter_has_zero_flag(
        self,
        encounter_df,
        hwr_excl_df,
        hwr_specialty_cohort_df,
        hwr_surg_gyn_df,
        empty_procedure_df,
    ) -> None:
        result = hwr_int_denominator(
            encounter_df,
            hwr_excl_df,
            hwr_specialty_cohort_df,
            hwr_surg_gyn_df,
            empty_procedure_df,
        )
        non_excl = result.filter(pl.col("exclusion_flag") == 0)
        assert len(non_excl) > 0


class TestHwrIntPlannedReadmission:
    """Tests for hwr_int_planned_readmission PAA classification."""

    @pytest.fixture
    def hwr_denominator_df(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "encounter_id": ["ENC001"],
                "person_id": ["P001"],
                "admission_date": [date(2025, 2, 1)],
                "discharge_date": [date(2025, 2, 5)],
                "discharge_disposition_code": ["01"],
                "facility_id": ["H001"],
                "principal_diagnosis_code": ["I50.1"],
                "ccs_diagnosis_category": ["108"],
                "specialty_cohort": ["CARDIORESPIRATORY"],
                "exclusion_flag": [0],
                "exclusion_reason": [None],
                "attributed_tin": [None],
                "attribution_role": [None],
            },
            schema={
                "encounter_id": pl.String,
                "person_id": pl.String,
                "admission_date": pl.Date,
                "discharge_date": pl.Date,
                "discharge_disposition_code": pl.String,
                "facility_id": pl.String,
                "principal_diagnosis_code": pl.String,
                "ccs_diagnosis_category": pl.String,
                "specialty_cohort": pl.String,
                "exclusion_flag": pl.Int32,
                "exclusion_reason": pl.String,
                "attributed_tin": pl.String,
                "attribution_role": pl.String,
            },
        )

    @pytest.fixture
    def hwr_readmission_encounter_df(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "encounter_id": ["ENC_HWR1", "ENC_HWR2"],
                "person_id": ["P001", "P001"],
                "encounter_type": ["acute inpatient", "acute inpatient"],
                "encounter_start_date": [date(2025, 2, 20), date(2025, 2, 25)],
                "encounter_end_date": [date(2025, 2, 23), date(2025, 2, 28)],
                "length_of_stay": [3, 3],
                "discharge_disposition_code": ["01", "01"],
                "facility_id": ["H001", "H001"],
                "primary_diagnosis_code": ["I21.0", "Z51.11"],
                "ccs_diagnosis_category": ["100", "45"],
                "drg_code_type": ["MS-DRG", "MS-DRG"],
                "drg_code": ["282", "834"],
                "encounter_group": ["claims", "claims"],
            },
            schema={
                "encounter_id": pl.String,
                "person_id": pl.String,
                "encounter_type": pl.String,
                "encounter_start_date": pl.Date,
                "encounter_end_date": pl.Date,
                "length_of_stay": pl.Int32,
                "discharge_disposition_code": pl.String,
                "facility_id": pl.String,
                "primary_diagnosis_code": pl.String,
                "ccs_diagnosis_category": pl.String,
                "drg_code_type": pl.String,
                "drg_code": pl.String,
                "encounter_group": pl.String,
            },
        )

    def test_returns_dataframe(
        self, hwr_denominator_df, hwr_readmission_encounter_df
    ) -> None:
        result = hwr_int_planned_readmission(
            hwr_readmission_encounter_df,
            hwr_denominator_df,
            _make_paa1_df(),
            _make_paa2_df(),
            _make_paa3_df(),
            _make_paa4_df(),
            _make_ccs_icd10_cm_df(),
        )
        assert isinstance(result, pl.DataFrame)

    def test_unplanned_readmission_detected(
        self, hwr_denominator_df, hwr_readmission_encounter_df
    ) -> None:
        result = hwr_int_planned_readmission(
            hwr_readmission_encounter_df,
            hwr_denominator_df,
            _make_paa1_df(),
            _make_paa2_df(),
            _make_paa3_df(),
            _make_paa4_df(),
            _make_ccs_icd10_cm_df(),
        )
        # ENC_HWR1 (AMI, 15 days after discharge) → unplanned
        hwra1 = result.filter(pl.col("readmission_encounter_id") == "ENC_HWR1")
        assert len(hwra1) == 1
        assert hwra1["unplanned_readmission_flag"][0] == 1

    def test_planned_readmission_not_in_numerator(
        self, hwr_denominator_df, hwr_readmission_encounter_df
    ) -> None:
        result = hwr_int_planned_readmission(
            hwr_readmission_encounter_df,
            hwr_denominator_df,
            _make_paa1_df(),
            _make_paa2_df(),
            _make_paa3_df(),
            _make_paa4_df(),
            _make_ccs_icd10_cm_df(),
        )
        # ENC_HWR2 (chemo/CCS 45 → Rule 2 planned) → unplanned_readmission_flag = 0
        hwra2 = result.filter(pl.col("readmission_encounter_id") == "ENC_HWR2")
        if len(hwra2) > 0:
            assert hwra2["unplanned_readmission_flag"][0] == 0

    def test_expected_columns(
        self, hwr_denominator_df, hwr_readmission_encounter_df
    ) -> None:
        result = hwr_int_planned_readmission(
            hwr_readmission_encounter_df,
            hwr_denominator_df,
            _make_paa1_df(),
            _make_paa2_df(),
            _make_paa3_df(),
            _make_paa4_df(),
            _make_ccs_icd10_cm_df(),
        )
        for col in [
            "index_encounter_id",
            "readmission_encounter_id",
            "person_id",
            "days_to_readmission",
            "is_within_30_days",
            "is_planned",
            "is_psychiatric_or_rehab",
            "unplanned_readmission_flag",
            "attributed_tin",
        ]:
            assert col in result.columns


# ═══════════════════════════════════════════════════════════════════════════════
# Pipeline tag tests
# ═══════════════════════════════════════════════════════════════════════════════


class TestPipelineTags:
    """Verify that the pipeline Tag references use the new measure/program namespaces."""

    def test_uamcc_refs_include_measure_tag(self) -> None:
        from aco.pipe.cms_quality_measures import pipeline_uamcc

        all_tags = [tag for expr in pipeline_uamcc.exprs for tag in expr.refs]
        tag_labels = [t.label for t in all_tags]
        assert "measure:UAMCC" in tag_labels

    def test_acr_refs_include_measure_tag(self) -> None:
        from aco.pipe.cms_quality_measures import pipeline_acr

        all_tags = [tag for expr in pipeline_acr.exprs for tag in expr.refs]
        tag_labels = [t.label for t in all_tags]
        assert "measure:ACR" in tag_labels

    def test_hwr_refs_include_measure_tag(self) -> None:
        from aco.pipe.cms_quality_measures import pipeline_hwr

        all_tags = [tag for expr in pipeline_hwr.exprs for tag in expr.refs]
        tag_labels = [t.label for t in all_tags]
        assert "measure:HWR" in tag_labels

    def test_reach_program_tag(self) -> None:
        from aco.pipe.cms_quality_measures import pipeline_uamcc

        all_tags = [tag for expr in pipeline_uamcc.exprs for tag in expr.refs]
        tag_labels = [t.label for t in all_tags]
        assert "program:reach" in tag_labels

    def test_mips_program_tag(self) -> None:
        from aco.pipe.cms_quality_measures import pipeline_hwr

        all_tags = [tag for expr in pipeline_hwr.exprs for tag in expr.refs]
        tag_labels = [t.label for t in all_tags]
        assert "program:mips" in tag_labels

    def test_pipeline_expr_names_are_qualified(self) -> None:
        from aco.pipe.cms_quality_measures import pipeline

        for expr in pipeline.exprs:
            assert "." in expr.name, f"{expr.name!r} must be qualified"

    def test_combined_pipeline_length(self) -> None:
        from aco.pipe.cms_quality_measures import (
            pipeline,
            pipeline_acr,
            pipeline_hwr,
            pipeline_uamcc,
        )

        assert len(pipeline.exprs) == (
            len(pipeline_uamcc.exprs)
            + len(pipeline_acr.exprs)
            + len(pipeline_hwr.exprs)
        )


# ═══════════════════════════════════════════════════════════════════════════
# uamcc_int_numerator
# ═══════════════════════════════════════════════════════════════════════════


class TestUamccIntNumerator:
    """uamcc_int_numerator filters claims to unplanned admissions in denom."""

    @pytest.fixture
    def stg_claims(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "claim_id": ["C001", "C002", "C003"],
                "person_id": ["P001", "P001", "P002"],
                "claim_start_date": [
                    date(2025, 3, 1),
                    date(2025, 5, 1),
                    date(2025, 4, 1),
                ],
                "claim_end_date": [
                    date(2025, 3, 5),
                    date(2025, 5, 5),
                    date(2025, 4, 5),
                ],
                "principal_diagnosis_code": ["I50.1", "E11.9", "J18.9"],
                "claim_type": ["acute inpatient", "acute inpatient", "acute inpatient"],
            }
        )

    @pytest.fixture
    def denominator(self) -> pl.DataFrame:
        return pl.DataFrame({"person_id": ["P001"]})

    @pytest.fixture
    def outcome_exclusion(self) -> pl.DataFrame:
        return pl.DataFrame({"claim_id": ["C002"]})  # exclude C002

    @pytest.fixture
    def value_set_ccs(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "icd_10_cm": ["I50.1"],
                "ccs_category": ["108"],
            }
        )

    def test_returns_dataframe(
        self, stg_claims, denominator, outcome_exclusion, value_set_ccs
    ):
        from aco.express.cms_quality_measures import uamcc_int_numerator

        result = uamcc_int_numerator(
            stg_claims, denominator, outcome_exclusion, value_set_ccs
        )
        assert isinstance(result, pl.DataFrame)

    def test_filters_to_denom_beneficiaries(
        self, stg_claims, denominator, outcome_exclusion, value_set_ccs
    ):
        from aco.express.cms_quality_measures import uamcc_int_numerator

        result = uamcc_int_numerator(
            stg_claims, denominator, outcome_exclusion, value_set_ccs
        )
        assert all(p == "P001" for p in result["person_id"].to_list())

    def test_excludes_outcome_exclusion_claims(
        self, stg_claims, denominator, outcome_exclusion, value_set_ccs
    ):
        from aco.express.cms_quality_measures import uamcc_int_numerator

        result = uamcc_int_numerator(
            stg_claims, denominator, outcome_exclusion, value_set_ccs
        )
        assert "C002" not in result["claim_id"].to_list()

    def test_expected_columns(
        self, stg_claims, denominator, outcome_exclusion, value_set_ccs
    ):
        from aco.express.cms_quality_measures import uamcc_int_numerator

        result = uamcc_int_numerator(
            stg_claims, denominator, outcome_exclusion, value_set_ccs
        )
        for col in (
            "person_id",
            "claim_id",
            "admission_date",
            "discharge_date",
            "principal_diagnosis_code",
            "unplanned_admission_flag",
        ):
            assert col in result.columns

    def test_unplanned_admission_flag_is_one(
        self, stg_claims, denominator, outcome_exclusion, value_set_ccs
    ):
        from aco.express.cms_quality_measures import uamcc_int_numerator

        result = uamcc_int_numerator(
            stg_claims, denominator, outcome_exclusion, value_set_ccs
        )
        assert all(f == 1 for f in result["unplanned_admission_flag"].to_list())


# ═══════════════════════════════════════════════════════════════════════════
# uamcc_summary
# ═══════════════════════════════════════════════════════════════════════════


class TestUamccSummary:
    """uamcc_summary computes observed UAMCC rate per 100 person-years."""

    @pytest.fixture
    def numerator_df(self) -> pl.DataFrame:
        return pl.DataFrame({"claim_id": ["C001", "C002", "C003"]})

    @pytest.fixture
    def person_time_df(self) -> pl.DataFrame:
        return pl.DataFrame({"person_id": ["P001", "P002"], "person_years": [0.8, 1.2]})

    @pytest.fixture
    def denominator_df(self) -> pl.DataFrame:
        return pl.DataFrame({"person_id": ["P001", "P002"]})

    def test_returns_dataframe(
        self, numerator_df, person_time_df, denominator_df, uamcc_period_df
    ):
        from aco.express.cms_quality_measures import uamcc_summary

        result = uamcc_summary(
            numerator_df, person_time_df, denominator_df, uamcc_period_df
        )
        assert isinstance(result, pl.DataFrame)

    def test_single_row(
        self, numerator_df, person_time_df, denominator_df, uamcc_period_df
    ):
        from aco.express.cms_quality_measures import uamcc_summary

        result = uamcc_summary(
            numerator_df, person_time_df, denominator_df, uamcc_period_df
        )
        assert len(result) == 1

    def test_performance_year(
        self, numerator_df, person_time_df, denominator_df, uamcc_period_df
    ):
        from aco.express.cms_quality_measures import uamcc_summary

        result = uamcc_summary(
            numerator_df, person_time_df, denominator_df, uamcc_period_df
        )
        assert result["performance_year"][0] == 2025

    def test_observed_admissions(
        self, numerator_df, person_time_df, denominator_df, uamcc_period_df
    ):
        from aco.express.cms_quality_measures import uamcc_summary

        result = uamcc_summary(
            numerator_df, person_time_df, denominator_df, uamcc_period_df
        )
        assert result["observed_admissions"][0] == 3

    def test_expected_columns(
        self, numerator_df, person_time_df, denominator_df, uamcc_period_df
    ):
        from aco.express.cms_quality_measures import uamcc_summary

        result = uamcc_summary(
            numerator_df, person_time_df, denominator_df, uamcc_period_df
        )
        for col in (
            "aco_id",
            "program",
            "performance_year",
            "denominator_count",
            "total_person_years",
            "observed_admissions",
            "observed_rate_per_100",
            "expected_admissions",
            "rsaar",
        ):
            assert col in result.columns


# ═══════════════════════════════════════════════════════════════════════════
# acr_summary
# ═══════════════════════════════════════════════════════════════════════════


class TestAcrSummary:
    """acr_summary computes observed ACR readmission rate."""

    @pytest.fixture
    def index_admission_df(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "encounter_id": ["E001", "E002", "E003"],
                "exclusion_flag": [0, 0, 1],  # 2 eligible
            }
        )

    @pytest.fixture
    def planned_readmission_df(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "readmission_encounter_id": ["R001"],
                "unplanned_readmission_flag": [1],
            }
        )

    def test_returns_dataframe(
        self, index_admission_df, planned_readmission_df, acr_period_df
    ):
        from aco.express.cms_quality_measures import acr_summary

        result = acr_summary(index_admission_df, planned_readmission_df, acr_period_df)
        assert isinstance(result, pl.DataFrame)

    def test_single_row(
        self, index_admission_df, planned_readmission_df, acr_period_df
    ):
        from aco.express.cms_quality_measures import acr_summary

        result = acr_summary(index_admission_df, planned_readmission_df, acr_period_df)
        assert len(result) == 1

    def test_denominator_count(
        self, index_admission_df, planned_readmission_df, acr_period_df
    ):
        from aco.express.cms_quality_measures import acr_summary

        result = acr_summary(index_admission_df, planned_readmission_df, acr_period_df)
        assert result["denominator_count"][0] == 2

    def test_observed_readmissions(
        self, index_admission_df, planned_readmission_df, acr_period_df
    ):
        from aco.express.cms_quality_measures import acr_summary

        result = acr_summary(index_admission_df, planned_readmission_df, acr_period_df)
        assert result["observed_readmissions"][0] == 1

    def test_expected_columns(
        self, index_admission_df, planned_readmission_df, acr_period_df
    ):
        from aco.express.cms_quality_measures import acr_summary

        result = acr_summary(index_admission_df, planned_readmission_df, acr_period_df)
        for col in (
            "aco_id",
            "program",
            "performance_year",
            "denominator_count",
            "observed_readmissions",
            "observed_rate",
            "expected_readmissions",
            "rsrr",
        ):
            assert col in result.columns


# ═══════════════════════════════════════════════════════════════════════════
# hwr_summary
# ═══════════════════════════════════════════════════════════════════════════


class TestHwrSummary:
    """hwr_summary computes observed HWR readmission rate."""

    @pytest.fixture
    def hwr_denom_df(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "encounter_id": ["E001", "E002", "E003"],
                "exclusion_flag": [0, 0, 0],
            }
        )

    @pytest.fixture
    def hwr_readmission_df(self) -> pl.DataFrame:
        return pl.DataFrame(
            {
                "readmission_encounter_id": ["R001", "R002"],
                "unplanned_readmission_flag": [1, 1],
            }
        )

    def test_returns_dataframe(self, hwr_denom_df, hwr_readmission_df, hwr_period_df):
        from aco.express.cms_quality_measures import hwr_summary

        result = hwr_summary(hwr_denom_df, hwr_readmission_df, hwr_period_df)
        assert isinstance(result, pl.DataFrame)

    def test_single_row(self, hwr_denom_df, hwr_readmission_df, hwr_period_df):
        from aco.express.cms_quality_measures import hwr_summary

        result = hwr_summary(hwr_denom_df, hwr_readmission_df, hwr_period_df)
        assert len(result) == 1

    def test_denominator_count(self, hwr_denom_df, hwr_readmission_df, hwr_period_df):
        from aco.express.cms_quality_measures import hwr_summary

        result = hwr_summary(hwr_denom_df, hwr_readmission_df, hwr_period_df)
        assert result["denominator_count"][0] == 3

    def test_observed_readmissions(
        self, hwr_denom_df, hwr_readmission_df, hwr_period_df
    ):
        from aco.express.cms_quality_measures import hwr_summary

        result = hwr_summary(hwr_denom_df, hwr_readmission_df, hwr_period_df)
        assert result["observed_readmissions"][0] == 2

    def test_expected_columns(self, hwr_denom_df, hwr_readmission_df, hwr_period_df):
        from aco.express.cms_quality_measures import hwr_summary

        result = hwr_summary(hwr_denom_df, hwr_readmission_df, hwr_period_df)
        for col in (
            "tin",
            "performance_year",
            "attribution_role",
            "denominator_count",
            "observed_readmissions",
            "observed_rate",
            "expected_readmissions",
            "rsrr",
        ):
            assert col in result.columns