stack/src/aco/express/ahrq_measures.py

from __future__ import annotations

import narwhals as nw
from narwhals.typing import FrameT

# ---------------------------------------------------------------------------
# Staging
# ---------------------------------------------------------------------------


@nw.narwhalify
def stg_pqi_inpatient_encounter(core__encounter: FrameT) -> FrameT:
    """Filter core.encounter to acute inpatient, add year_number."""
    return (
        core__encounter.filter(nw.col("encounter_type") == "acute inpatient")
        .with_columns(
            nw.col("encounter_start_date")
            .dt.year()
            .cast(nw.Int64)
            .alias("year_number"),
        )
        .select(
            "encounter_id",
            "data_source",
            "drg_code",
            "drg_description",
            "admit_source_code",
            "encounter_start_date",
            "encounter_end_date",
            "length_of_stay",
            "primary_diagnosis_code",
            "person_id",
            "facility_id",
            "paid_amount",
            "year_number",
        )
    )


@nw.narwhalify
def stg_pqi_member_months(
    core__member_months: FrameT,
    reference_data__calendar: FrameT,
) -> FrameT:
    """Join member_months with calendar to get first_day_of_month."""
    cal = (
        reference_data__calendar.with_columns(
            nw.col("year_month").str.replace_all("-", "").alias("yyyymm"),
        )
        .select("yyyymm", "first_day_of_month")
        .unique()
    )
    return core__member_months.join(
        cal,
        left_on=["year_month"],
        right_on=["yyyymm"],
        how="inner",
    ).select(
        "data_source",
        "person_id",
        "first_day_of_month",
        nw.col("year_month"),
    )


# ---------------------------------------------------------------------------
# Shared exclusions
# ---------------------------------------------------------------------------


@nw.narwhalify
def int_pqi_shared_exclusion_missing_dates(
    ahrq_measures___stg_pqi_inpatient_encounter: FrameT,
) -> FrameT:
    """Encounters with missing start date."""
    return ahrq_measures___stg_pqi_inpatient_encounter.filter(
        nw.col("encounter_start_date").is_null()
    ).select("encounter_id", "data_source")


@nw.narwhalify
def int_pqi_shared_exclusion_missing_gender(
    core__patient: FrameT,
) -> FrameT:
    """Patients with sex not in (male, female)."""
    return core__patient.filter(~nw.col("sex").is_in(["male", "female"])).select(
        "data_source", "person_id"
    )


@nw.narwhalify
def int_pqi_shared_exclusion_missing_primary_dx(
    ahrq_measures___stg_pqi_inpatient_encounter: FrameT,
) -> FrameT:
    """Encounters with missing primary diagnosis."""
    return ahrq_measures___stg_pqi_inpatient_encounter.filter(
        nw.col("primary_diagnosis_code").is_null()
    ).select("encounter_id", "data_source")


@nw.narwhalify
def int_pqi_shared_exclusion_transfer(
    ahrq_measures___stg_pqi_inpatient_encounter: FrameT,
) -> FrameT:
    """Encounters admitted via transfer."""
    return ahrq_measures___stg_pqi_inpatient_encounter.filter(
        nw.col("admit_source_code").is_in(["4", "5", "6"])
    ).select("encounter_id", "data_source")


@nw.narwhalify
def int_pqi_shared_exclusion_ungroupable_drg(
    ahrq_measures___stg_pqi_inpatient_encounter: FrameT,
) -> FrameT:
    """Encounters with ungroupable DRG (999)."""
    return ahrq_measures___stg_pqi_inpatient_encounter.filter(
        nw.col("drg_code") == "999"
    ).select("encounter_id", "data_source")


@nw.narwhalify
def int_pqi_shared_exclusions_missing_age(
    core__patient: FrameT,
) -> FrameT:
    """Patients with missing birth_date."""
    return core__patient.filter(nw.col("birth_date").is_null()).select(
        "data_source", "person_id"
    )


@nw.narwhalify
def int_pqi_shared_exclusion_union(
    ahrq_measures___stg_pqi_inpatient_encounter: FrameT,
    ahrq_measures___int_pqi_shared_exclusions_missing_age: FrameT,
    ahrq_measures___int_pqi_shared_exclusion_missing_gender: FrameT,
    ahrq_measures___int_pqi_shared_exclusion_missing_dates: FrameT,
    ahrq_measures___int_pqi_shared_exclusion_missing_primary_dx: FrameT,
    ahrq_measures___int_pqi_shared_exclusion_transfer: FrameT,
    ahrq_measures___int_pqi_shared_exclusion_ungroupable_drg: FrameT,
) -> FrameT:
    """Union all shared exclusions, joining person-level ones back
    through _stg_pqi_inpatient_encounter to get encounter_id."""
    denom = ahrq_measures___stg_pqi_inpatient_encounter

    # Person-level exclusions → join to encounters
    age_exc = (
        denom.join(
            ahrq_measures___int_pqi_shared_exclusions_missing_age,
            on=["person_id", "data_source"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .with_columns(nw.lit("missing age").alias("exclusion_reason"))
    )
    gender_exc = (
        denom.join(
            ahrq_measures___int_pqi_shared_exclusion_missing_gender,
            on=["person_id", "data_source"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .with_columns(nw.lit("missing gender").alias("exclusion_reason"))
    )

    # Encounter-level exclusions
    dates_exc = (
        denom.join(
            ahrq_measures___int_pqi_shared_exclusion_missing_dates,
            on=["encounter_id", "data_source"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .with_columns(nw.lit("missing dates").alias("exclusion_reason"))
    )
    dx_exc = (
        denom.join(
            ahrq_measures___int_pqi_shared_exclusion_missing_primary_dx,
            on=["encounter_id", "data_source"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .with_columns(nw.lit("missing primary diagnosis").alias("exclusion_reason"))
    )
    tx_exc = (
        denom.join(
            ahrq_measures___int_pqi_shared_exclusion_transfer,
            on=["encounter_id", "data_source"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .with_columns(nw.lit("transfer").alias("exclusion_reason"))
    )
    drg_exc = (
        denom.join(
            ahrq_measures___int_pqi_shared_exclusion_ungroupable_drg,
            on=["encounter_id", "data_source"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .with_columns(nw.lit("ungroupable DRG").alias("exclusion_reason"))
    )

    return nw.concat([age_exc, gender_exc, dates_exc, dx_exc, tx_exc, drg_exc])


# ---------------------------------------------------------------------------
# Helpers for repeated patterns
# ---------------------------------------------------------------------------


def _pqi_denom(
    mm: FrameT,
    patient: FrameT,
    min_age: int = 18,
    max_age: int | None = None,
) -> FrameT:
    """Common denom logic: join member_months to patient, filter by age.

    SQL: date_diff('year', birth_date, first_day_of_month) >= min_age
    """
    result = mm.join(
        patient.select("person_id", "data_source", "birth_date"),
        on=["person_id", "data_source"],
        how="inner",
    )
    age_expr = (
        (
            nw.col("first_day_of_month").cast(nw.Datetime)
            - nw.col("birth_date").cast(nw.Datetime)
        ).dt.total_seconds()
        / 86400
        / 365
    ).cast(nw.Int64)

    cond = age_expr >= min_age
    if max_age is not None:
        cond = cond & (age_expr <= max_age)

    result = result.filter(cond)
    return result.select(
        nw.col("year_month").str.slice(0, 4).cast(nw.Int64).alias("year_number"),
        "person_id",
        "data_source",
        age_expr.alias("age"),
    ).unique()


def _pqi_simple_exclusions(
    shared_union: FrameT,
) -> FrameT:
    """PQI measures that use only the shared exclusion union (01, 03, 14)."""
    return shared_union.with_columns(
        nw.col("exclusion_reason")
        .rank(method="ordinal")
        .over("encounter_id", "data_source")
        .cast(nw.Int64)
        .alias("exclusion_number"),
    ).select(
        "encounter_id",
        "data_source",
        "exclusion_reason",
        "exclusion_number",
    )


def _pqi_num(
    encounters: FrameT,
    value_set: FrameT,
    denom: FrameT,
    exclusions: FrameT,
    value_set_name: str,
    pqi_number: str,
) -> FrameT:
    """Common numerator: inpatient encounters matching a value set,
    in denom, not excluded."""
    pqi_codes = (
        value_set.filter(
            (nw.col("value_set_name") == value_set_name)
            & (nw.col("pqi_number") == pqi_number)
        )
        .select(nw.col("code"))
        .unique()
    )
    result = encounters.join(
        pqi_codes,
        left_on=["primary_diagnosis_code"],
        right_on=["code"],
        how="inner",
    )
    result = result.join(
        denom.select("person_id", "data_source", "year_number"),
        on=["person_id", "data_source", "year_number"],
        how="inner",
    )
    excl_marker = (
        exclusions.select("encounter_id", "data_source")
        .unique()
        .with_columns(nw.lit(1).alias("_excl"))
    )
    result = result.join(
        excl_marker,
        on=["encounter_id", "data_source"],
        how="left",
    )
    return result.filter(nw.col("_excl").is_null()).select(
        "data_source",
        "person_id",
        "year_number",
        "encounter_id",
    )


# ---------------------------------------------------------------------------
# PQI 01 — Diabetes short-term complications
# ---------------------------------------------------------------------------


@nw.narwhalify
def int_pqi_01_denom(
    ahrq_measures___stg_pqi_member_months: FrameT,
    core__patient: FrameT,
) -> FrameT:
    """Eligible members for PQI 01 (diabetes short-term complications), all ages."""
    return _pqi_denom(ahrq_measures___stg_pqi_member_months, core__patient)


@nw.narwhalify
def int_pqi_01_exclusions(
    ahrq_measures___int_pqi_shared_exclusion_union: FrameT,
) -> FrameT:
    """Shared exclusions only (missing data, transfer, ungroupable DRG)."""
    return _pqi_simple_exclusions(ahrq_measures___int_pqi_shared_exclusion_union)


@nw.narwhalify
def int_pqi_01_num(
    ahrq_measures___stg_pqi_inpatient_encounter: FrameT,
    ahrq_measures___value_set_pqi: FrameT,
    ahrq_measures___int_pqi_01_denom: FrameT,
    ahrq_measures___int_pqi_01_exclusions: FrameT,
) -> FrameT:
    """Diabetes short-term complications diagnosis, in denom, not excluded."""
    return _pqi_num(
        ahrq_measures___stg_pqi_inpatient_encounter,
        ahrq_measures___value_set_pqi,
        ahrq_measures___int_pqi_01_denom,
        ahrq_measures___int_pqi_01_exclusions,
        "diabetes_with_short-term_complications_diagnosis_codes",
        "01",
    )


# ---------------------------------------------------------------------------
# PQI 03 — Diabetes long-term complications
# ---------------------------------------------------------------------------


@nw.narwhalify
def int_pqi_03_denom(
    ahrq_measures___stg_pqi_member_months: FrameT,
    core__patient: FrameT,
) -> FrameT:
    """Eligible members for PQI 03 (diabetes long-term complications), all ages."""
    return _pqi_denom(ahrq_measures___stg_pqi_member_months, core__patient)


@nw.narwhalify
def int_pqi_03_exclusions(
    ahrq_measures___int_pqi_shared_exclusion_union: FrameT,
) -> FrameT:
    """Shared exclusions only (missing data, transfer, ungroupable DRG)."""
    return _pqi_simple_exclusions(ahrq_measures___int_pqi_shared_exclusion_union)


@nw.narwhalify
def int_pqi_03_num(
    ahrq_measures___stg_pqi_inpatient_encounter: FrameT,
    ahrq_measures___value_set_pqi: FrameT,
    ahrq_measures___int_pqi_03_denom: FrameT,
    ahrq_measures___int_pqi_03_exclusions: FrameT,
) -> FrameT:
    """Diabetes long-term complications diagnosis, in denom, not excluded."""
    return _pqi_num(
        ahrq_measures___stg_pqi_inpatient_encounter,
        ahrq_measures___value_set_pqi,
        ahrq_measures___int_pqi_03_denom,
        ahrq_measures___int_pqi_03_exclusions,
        "diabetes_with_long-term_complications_diagnosis_codes",
        "03",
    )


# ---------------------------------------------------------------------------
# PQI 05 — COPD/Asthma (age >= 40, cystic fibrosis exclusion)
# ---------------------------------------------------------------------------


@nw.narwhalify
def int_pqi_05_denom(
    ahrq_measures___stg_pqi_member_months: FrameT,
    core__patient: FrameT,
) -> FrameT:
    """Eligible members for PQI 05 (COPD/asthma), age >= 40."""
    return _pqi_denom(
        ahrq_measures___stg_pqi_member_months,
        core__patient,
        min_age=40,
    )


@nw.narwhalify
def int_pqi_05_exclusions(
    ahrq_measures___int_pqi_shared_exclusion_union: FrameT,
    ahrq_measures___value_set_pqi: FrameT,
    core__condition: FrameT,
) -> FrameT:
    """Shared exclusions + cystic fibrosis from conditions."""
    cf_codes = (
        ahrq_measures___value_set_pqi.filter(
            (
                nw.col("value_set_name")
                == "cystic_fibrosis_and_anomalies_of_the_respiratory_system"
            )
            & (nw.col("pqi_number") == "05")
        )
        .select("code")
        .unique()
    )
    resp_an = (
        core__condition.filter(nw.col("normalized_code_type") == "icd-10-cm")
        .filter(~nw.col("encounter_id").is_null())
        .join(
            cf_codes,
            left_on=["normalized_code"],
            right_on=["code"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .unique()
        .with_columns(nw.lit("cystic fibrosis").alias("exclusion_reason"))
    )
    combined = nw.concat(
        [
            ahrq_measures___int_pqi_shared_exclusion_union.select(
                "encounter_id", "data_source", "exclusion_reason"
            ),
            resp_an,
        ]
    )
    return combined.with_columns(
        nw.col("exclusion_reason")
        .rank(method="ordinal")
        .over("encounter_id", "data_source")
        .cast(nw.Int64)
        .alias("exclusion_number"),
    ).select(
        "encounter_id",
        "data_source",
        "exclusion_reason",
        "exclusion_number",
    )


@nw.narwhalify
def int_pqi_05_num(
    ahrq_measures___stg_pqi_inpatient_encounter: FrameT,
    ahrq_measures___value_set_pqi: FrameT,
    ahrq_measures___int_pqi_05_denom: FrameT,
    ahrq_measures___int_pqi_05_exclusions: FrameT,
) -> FrameT:
    """COPD or asthma diagnosis, in denom, not excluded."""
    e = ahrq_measures___stg_pqi_inpatient_encounter
    vs = ahrq_measures___value_set_pqi

    copd_codes = (
        vs.filter(
            (nw.col("value_set_name") == "chronic_obstructive_pulmonary_disorder")
            & (nw.col("pqi_number") == "05")
        )
        .select("code")
        .unique()
    )
    asthma_codes = (
        vs.filter(
            (nw.col("value_set_name") == "asthma") & (nw.col("pqi_number") == "05")
        )
        .select("code")
        .unique()
    )

    copd_marker = copd_codes.with_columns(nw.lit(1).alias("_copd"))
    asthma_marker = asthma_codes.with_columns(nw.lit(1).alias("_asthma"))

    result = e.join(
        copd_marker,
        left_on=["primary_diagnosis_code"],
        right_on=["code"],
        how="left",
    ).join(
        asthma_marker,
        left_on=["primary_diagnosis_code"],
        right_on=["code"],
        how="left",
    )
    result = result.filter(~nw.col("_copd").is_null() | ~nw.col("_asthma").is_null())
    result = result.join(
        ahrq_measures___int_pqi_05_denom.select(
            "person_id", "data_source", "year_number"
        ),
        on=["person_id", "data_source", "year_number"],
        how="inner",
    )
    excl = (
        ahrq_measures___int_pqi_05_exclusions.select("encounter_id", "data_source")
        .unique()
        .with_columns(nw.lit(1).alias("_excl"))
    )
    result = result.join(excl, on=["encounter_id", "data_source"], how="left")
    return result.filter(nw.col("_excl").is_null()).select(
        "data_source", "person_id", "year_number", "encounter_id"
    )


# ---------------------------------------------------------------------------
# PQI 07 — Hypertension (cardiac procedure + CKD exclusions)
# ---------------------------------------------------------------------------


@nw.narwhalify
def int_pqi_07_denom(
    ahrq_measures___stg_pqi_member_months: FrameT,
    core__patient: FrameT,
) -> FrameT:
    """Eligible members for PQI 07 (hypertension), all ages."""
    return _pqi_denom(ahrq_measures___stg_pqi_member_months, core__patient)


@nw.narwhalify
def int_pqi_07_exclusions(
    ahrq_measures___int_pqi_shared_exclusion_union: FrameT,
    ahrq_measures___value_set_pqi: FrameT,
    core__condition: FrameT,
    core__procedure: FrameT,
) -> FrameT:
    """Shared + cardiac procedure + CKD w/ dialysis access."""
    vs = ahrq_measures___value_set_pqi
    # Cardiac procedure
    cardiac_codes = (
        vs.filter(
            (nw.col("value_set_name") == "cardiac_procedure_codes")
            & (nw.col("pqi_number") == "appendix_b")
        )
        .select("code")
        .unique()
    )
    cardiac = (
        core__procedure.filter(nw.col("normalized_code_type") == "icd-10-pcs")
        .filter(~nw.col("encounter_id").is_null())
        .join(
            cardiac_codes,
            left_on=["normalized_code"],
            right_on=["code"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .unique()
        .with_columns(nw.lit("cardiac procedure").alias("exclusion_reason"))
    )
    # CKD
    ckd_codes = (
        vs.filter(
            (nw.col("value_set_name") == "exclusion_kidney_disease_diagnosis_codes")
            & (nw.col("pqi_number") == "07")
        )
        .select("code")
        .unique()
    )
    ckd = (
        core__condition.filter(nw.col("normalized_code_type") == "icd-10-cm")
        .filter(~nw.col("encounter_id").is_null())
        .join(
            ckd_codes,
            left_on=["normalized_code"],
            right_on=["code"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .unique()
    )
    # Dialysis access
    access_codes = (
        vs.filter(
            (nw.col("value_set_name") == "exclusion_dialysis_access_procedure_codes")
            & (nw.col("pqi_number") == "07")
        )
        .select("code")
        .unique()
    )
    access_enc = (
        core__procedure.filter(nw.col("normalized_code_type") == "icd-10-pcs")
        .filter(~nw.col("encounter_id").is_null())
        .join(
            access_codes,
            left_on=["normalized_code"],
            right_on=["code"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .unique()
    )
    # CKD + dialysis access combined
    ckd_combined = (
        ckd.join(
            access_enc,
            on=["encounter_id", "data_source"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .with_columns(nw.lit("ckd").alias("exclusion_reason"))
    )

    combined = nw.concat(
        [
            ahrq_measures___int_pqi_shared_exclusion_union.select(
                "encounter_id", "data_source", "exclusion_reason"
            ),
            cardiac,
            ckd_combined,
        ]
    )
    return combined.with_columns(
        nw.col("exclusion_reason")
        .rank(method="ordinal")
        .over("encounter_id", "data_source")
        .cast(nw.Int64)
        .alias("exclusion_number"),
    ).select(
        "encounter_id",
        "data_source",
        "exclusion_reason",
        "exclusion_number",
    )


@nw.narwhalify
def int_pqi_07_num(
    ahrq_measures___stg_pqi_inpatient_encounter: FrameT,
    ahrq_measures___value_set_pqi: FrameT,
    ahrq_measures___int_pqi_07_denom: FrameT,
    ahrq_measures___int_pqi_07_exclusions: FrameT,
) -> FrameT:
    """Hypertension diagnosis, in denom, not excluded."""
    return _pqi_num(
        ahrq_measures___stg_pqi_inpatient_encounter,
        ahrq_measures___value_set_pqi,
        ahrq_measures___int_pqi_07_denom,
        ahrq_measures___int_pqi_07_exclusions,
        "hypertension_diagnosis_codes",
        "07",
    )


# ---------------------------------------------------------------------------
# PQI 08 — Heart failure (cardiac procedure exclusion)
# ---------------------------------------------------------------------------


@nw.narwhalify
def int_pqi_08_denom(
    ahrq_measures___stg_pqi_member_months: FrameT,
    core__patient: FrameT,
) -> FrameT:
    """Eligible members for PQI 08 (heart failure), all ages."""
    return _pqi_denom(ahrq_measures___stg_pqi_member_months, core__patient)


@nw.narwhalify
def int_pqi_08_exclusions(
    ahrq_measures___int_pqi_shared_exclusion_union: FrameT,
    ahrq_measures___value_set_pqi: FrameT,
    core__procedure: FrameT,
) -> FrameT:
    """Shared + cardiac procedure."""
    cardiac_codes = (
        ahrq_measures___value_set_pqi.filter(
            (nw.col("value_set_name") == "cardiac_procedure_codes")
            & (nw.col("pqi_number") == "appendix_b")
        )
        .select("code")
        .unique()
    )
    cardiac = (
        core__procedure.filter(nw.col("normalized_code_type") == "icd-10-pcs")
        .filter(~nw.col("encounter_id").is_null())
        .join(
            cardiac_codes,
            left_on=["normalized_code"],
            right_on=["code"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .unique()
        .with_columns(nw.lit("cardiac procedure").alias("exclusion_reason"))
    )
    combined = nw.concat(
        [
            ahrq_measures___int_pqi_shared_exclusion_union.select(
                "encounter_id", "data_source", "exclusion_reason"
            ),
            cardiac,
        ]
    )
    return combined.with_columns(
        nw.col("exclusion_reason")
        .rank(method="ordinal")
        .over("encounter_id", "data_source")
        .cast(nw.Int64)
        .alias("exclusion_number"),
    ).select(
        "encounter_id",
        "data_source",
        "exclusion_reason",
        "exclusion_number",
    )


@nw.narwhalify
def int_pqi_08_num(
    ahrq_measures___stg_pqi_inpatient_encounter: FrameT,
    ahrq_measures___value_set_pqi: FrameT,
    ahrq_measures___int_pqi_08_denom: FrameT,
    ahrq_measures___int_pqi_08_exclusions: FrameT,
) -> FrameT:
    """Heart failure diagnosis, in denom, not excluded."""
    return _pqi_num(
        ahrq_measures___stg_pqi_inpatient_encounter,
        ahrq_measures___value_set_pqi,
        ahrq_measures___int_pqi_08_denom,
        ahrq_measures___int_pqi_08_exclusions,
        "heart_failure_diagnosis_codes",
        "08",
    )


# ---------------------------------------------------------------------------
# PQI 11 — Bacterial pneumonia (sickle cell + immunocompromised)
# ---------------------------------------------------------------------------


@nw.narwhalify
def int_pqi_11_denom(
    ahrq_measures___stg_pqi_member_months: FrameT,
    core__patient: FrameT,
) -> FrameT:
    """Eligible members for PQI 11 (bacterial pneumonia), all ages."""
    return _pqi_denom(ahrq_measures___stg_pqi_member_months, core__patient)


@nw.narwhalify
def int_pqi_11_exclusions(
    ahrq_measures___int_pqi_shared_exclusion_union: FrameT,
    ahrq_measures___value_set_pqi: FrameT,
    core__condition: FrameT,
    core__procedure: FrameT,
) -> FrameT:
    """Shared + sickle cell + immunocompromised dx/px."""
    vs = ahrq_measures___value_set_pqi

    sickle_codes = (
        vs.filter(
            (
                nw.col("value_set_name")
                == "sickle_cell_anemia_or_hb-s_disease_diagnosis_codes"
            )
            & (nw.col("pqi_number") == "11")
        )
        .select("code")
        .unique()
    )
    sickle = (
        core__condition.filter(nw.col("normalized_code_type") == "icd-10-cm")
        .filter(~nw.col("encounter_id").is_null())
        .join(
            sickle_codes,
            left_on=["normalized_code"],
            right_on=["code"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .unique()
        .with_columns(nw.lit("sickle cell").alias("exclusion_reason"))
    )

    immune_dx_codes = (
        vs.filter(
            (nw.col("value_set_name") == "immunocompromised_state_diagnosis_codes")
            & (nw.col("pqi_number") == "appendix_c")
        )
        .select("code")
        .unique()
    )
    immune_dx = (
        core__condition.filter(nw.col("normalized_code_type") == "icd-10-cm")
        .filter(~nw.col("encounter_id").is_null())
        .join(
            immune_dx_codes,
            left_on=["normalized_code"],
            right_on=["code"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .unique()
        .with_columns(nw.lit("immunocompromised diagnosis").alias("exclusion_reason"))
    )

    immune_px_codes = (
        vs.filter(
            (nw.col("value_set_name") == "immunocompromised_state_procedure_codes")
            & (nw.col("pqi_number") == "appendix_c")
        )
        .select("code")
        .unique()
    )
    immune_px = (
        core__procedure.filter(nw.col("normalized_code_type") == "icd-10-pcs")
        .filter(~nw.col("encounter_id").is_null())
        .join(
            immune_px_codes,
            left_on=["normalized_code"],
            right_on=["code"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .unique()
        .with_columns(nw.lit("immunocompromised procedure").alias("exclusion_reason"))
    )

    combined = nw.concat(
        [
            ahrq_measures___int_pqi_shared_exclusion_union.select(
                "encounter_id", "data_source", "exclusion_reason"
            ),
            sickle,
            immune_dx,
            immune_px,
        ]
    )
    return combined.with_columns(
        nw.col("exclusion_reason")
        .rank(method="ordinal")
        .over("encounter_id", "data_source")
        .cast(nw.Int64)
        .alias("exclusion_number"),
    ).select(
        "encounter_id",
        "data_source",
        "exclusion_reason",
        "exclusion_number",
    )


@nw.narwhalify
def int_pqi_11_num(
    ahrq_measures___stg_pqi_inpatient_encounter: FrameT,
    ahrq_measures___value_set_pqi: FrameT,
    ahrq_measures___int_pqi_11_denom: FrameT,
    ahrq_measures___int_pqi_11_exclusions: FrameT,
) -> FrameT:
    """Community-acquired bacterial pneumonia diagnosis, in denom, not excluded."""
    return _pqi_num(
        ahrq_measures___stg_pqi_inpatient_encounter,
        ahrq_measures___value_set_pqi,
        ahrq_measures___int_pqi_11_denom,
        ahrq_measures___int_pqi_11_exclusions,
        "community_acquired_bacterial_pneumonia_diagnosis_codes",
        "11",
    )


# ---------------------------------------------------------------------------
# PQI 12 — UTI (kidney + immunocompromised)
# ---------------------------------------------------------------------------


@nw.narwhalify
def int_pqi_12_denom(
    ahrq_measures___stg_pqi_member_months: FrameT,
    core__patient: FrameT,
) -> FrameT:
    """Eligible members for PQI 12 (urinary tract infection), all ages."""
    return _pqi_denom(ahrq_measures___stg_pqi_member_months, core__patient)


@nw.narwhalify
def int_pqi_12_exclusions(
    ahrq_measures___int_pqi_shared_exclusion_union: FrameT,
    ahrq_measures___value_set_pqi: FrameT,
    core__condition: FrameT,
    core__procedure: FrameT,
) -> FrameT:
    """Shared + kidney + immunocompromised dx/px."""
    vs = ahrq_measures___value_set_pqi

    kidney_codes = (
        vs.filter(
            (
                nw.col("value_set_name")
                == "kidney_or_urinary_tract_disorder_diagnosis_codes"
            )
            & (nw.col("pqi_number") == "12")
        )
        .select("code")
        .unique()
    )
    kidney = (
        core__condition.filter(nw.col("normalized_code_type") == "icd-10-cm")
        .filter(~nw.col("encounter_id").is_null())
        .join(
            kidney_codes,
            left_on=["normalized_code"],
            right_on=["code"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .unique()
        .with_columns(nw.lit("kidney").alias("exclusion_reason"))
    )

    immune_dx_codes = (
        vs.filter(
            (nw.col("value_set_name") == "immunocompromised_state_diagnosis_codes")
            & (nw.col("pqi_number") == "appendix_c")
        )
        .select("code")
        .unique()
    )
    immune_dx = (
        core__condition.filter(nw.col("normalized_code_type") == "icd-10-cm")
        .filter(~nw.col("encounter_id").is_null())
        .join(
            immune_dx_codes,
            left_on=["normalized_code"],
            right_on=["code"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .unique()
        .with_columns(nw.lit("immunocompromised diagnosis").alias("exclusion_reason"))
    )

    immune_px_codes = (
        vs.filter(
            (nw.col("value_set_name") == "immunocompromised_state_procedure_codes")
            & (nw.col("pqi_number") == "appendix_c")
        )
        .select("code")
        .unique()
    )
    immune_px = (
        core__procedure.filter(nw.col("normalized_code_type") == "icd-10-pcs")
        .filter(~nw.col("encounter_id").is_null())
        .join(
            immune_px_codes,
            left_on=["normalized_code"],
            right_on=["code"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .unique()
        .with_columns(nw.lit("immunocompromised procedure").alias("exclusion_reason"))
    )

    combined = nw.concat(
        [
            ahrq_measures___int_pqi_shared_exclusion_union.select(
                "encounter_id", "data_source", "exclusion_reason"
            ),
            kidney,
            immune_dx,
            immune_px,
        ]
    )
    return combined.with_columns(
        nw.col("exclusion_reason")
        .rank(method="ordinal")
        .over("encounter_id", "data_source")
        .cast(nw.Int64)
        .alias("exclusion_number"),
    ).select(
        "encounter_id",
        "data_source",
        "exclusion_reason",
        "exclusion_number",
    )


@nw.narwhalify
def int_pqi_12_num(
    ahrq_measures___stg_pqi_inpatient_encounter: FrameT,
    ahrq_measures___value_set_pqi: FrameT,
    ahrq_measures___int_pqi_12_denom: FrameT,
    ahrq_measures___int_pqi_12_exclusions: FrameT,
) -> FrameT:
    """Urinary tract infection diagnosis, in denom, not excluded."""
    return _pqi_num(
        ahrq_measures___stg_pqi_inpatient_encounter,
        ahrq_measures___value_set_pqi,
        ahrq_measures___int_pqi_12_denom,
        ahrq_measures___int_pqi_12_exclusions,
        "urinary_tract_infection_diagnosis_codes",
        "12",
    )


# ---------------------------------------------------------------------------
# PQI 14 — Uncontrolled diabetes (simple exclusions)
# ---------------------------------------------------------------------------


@nw.narwhalify
def int_pqi_14_denom(
    ahrq_measures___stg_pqi_member_months: FrameT,
    core__patient: FrameT,
) -> FrameT:
    """Eligible members for PQI 14 (uncontrolled diabetes), all ages."""
    return _pqi_denom(ahrq_measures___stg_pqi_member_months, core__patient)


@nw.narwhalify
def int_pqi_14_exclusions(
    ahrq_measures___int_pqi_shared_exclusion_union: FrameT,
) -> FrameT:
    """Shared exclusions only (missing data, transfer, ungroupable DRG)."""
    return _pqi_simple_exclusions(ahrq_measures___int_pqi_shared_exclusion_union)


@nw.narwhalify
def int_pqi_14_num(
    ahrq_measures___stg_pqi_inpatient_encounter: FrameT,
    ahrq_measures___value_set_pqi: FrameT,
    ahrq_measures___int_pqi_14_denom: FrameT,
    ahrq_measures___int_pqi_14_exclusions: FrameT,
) -> FrameT:
    """Uncontrolled diabetes diagnosis, in denom, not excluded."""
    return _pqi_num(
        ahrq_measures___stg_pqi_inpatient_encounter,
        ahrq_measures___value_set_pqi,
        ahrq_measures___int_pqi_14_denom,
        ahrq_measures___int_pqi_14_exclusions,
        "uncontrolled_diabetes_diagnosis_codes",
        "14",
    )


# ---------------------------------------------------------------------------
# PQI 15 — Asthma, younger adults (18-39, cystic fibrosis)
# ---------------------------------------------------------------------------


@nw.narwhalify
def int_pqi_15_denom(
    ahrq_measures___stg_pqi_member_months: FrameT,
    core__patient: FrameT,
) -> FrameT:
    """Eligible members for PQI 15 (asthma, younger adults), age 18–39."""
    return _pqi_denom(
        ahrq_measures___stg_pqi_member_months,
        core__patient,
        min_age=18,
        max_age=39,
    )


@nw.narwhalify
def int_pqi_15_exclusions(
    ahrq_measures___int_pqi_shared_exclusion_union: FrameT,
    ahrq_measures___value_set_pqi: FrameT,
    core__condition: FrameT,
) -> FrameT:
    """Shared + cystic fibrosis."""
    cf_codes = (
        ahrq_measures___value_set_pqi.filter(
            (
                nw.col("value_set_name")
                == "cystic_fibrosis_and_anomalies_of_the_respiratory_system_diagnosis_codes"
            )
            & (nw.col("pqi_number") == "15")
        )
        .select("code")
        .unique()
    )
    resp_an = (
        core__condition.filter(nw.col("normalized_code_type") == "icd-10-cm")
        .filter(~nw.col("encounter_id").is_null())
        .join(
            cf_codes,
            left_on=["normalized_code"],
            right_on=["code"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .unique()
        .with_columns(nw.lit("cystic fibrosis").alias("exclusion_reason"))
    )
    combined = nw.concat(
        [
            ahrq_measures___int_pqi_shared_exclusion_union.select(
                "encounter_id", "data_source", "exclusion_reason"
            ),
            resp_an,
        ]
    )
    return combined.with_columns(
        nw.col("exclusion_reason")
        .rank(method="ordinal")
        .over("encounter_id", "data_source")
        .cast(nw.Int64)
        .alias("exclusion_number"),
    ).select(
        "encounter_id",
        "data_source",
        "exclusion_reason",
        "exclusion_number",
    )


@nw.narwhalify
def int_pqi_15_num(
    ahrq_measures___stg_pqi_inpatient_encounter: FrameT,
    ahrq_measures___value_set_pqi: FrameT,
    ahrq_measures___int_pqi_15_denom: FrameT,
    ahrq_measures___int_pqi_15_exclusions: FrameT,
) -> FrameT:
    """Asthma diagnosis (age 18–39), in denom, not excluded."""
    return _pqi_num(
        ahrq_measures___stg_pqi_inpatient_encounter,
        ahrq_measures___value_set_pqi,
        ahrq_measures___int_pqi_15_denom,
        ahrq_measures___int_pqi_15_exclusions,
        "asthma_diagnosis_codes",
        "15",
    )


# ---------------------------------------------------------------------------
# PQI 16 — Lower-extremity amputation (amputation + pregnancy)
# ---------------------------------------------------------------------------


@nw.narwhalify
def int_pqi_16_denom(
    ahrq_measures___stg_pqi_member_months: FrameT,
    core__patient: FrameT,
) -> FrameT:
    """Eligible members for PQI 16 (lower-extremity amputation), all ages."""
    return _pqi_denom(ahrq_measures___stg_pqi_member_months, core__patient)


@nw.narwhalify
def int_pqi_16_exclusions(
    ahrq_measures___int_pqi_shared_exclusion_union: FrameT,
    ahrq_measures___value_set_pqi: FrameT,
    core__condition: FrameT,
) -> FrameT:
    """Shared + amputation + pregnancy."""
    vs = ahrq_measures___value_set_pqi

    amp_codes = (
        vs.filter(
            (
                nw.col("value_set_name")
                == "traumatic_amputation_of_the_lower_extremity_diagnosis_codes"
            )
            & (nw.col("pqi_number") == "16")
        )
        .select("code")
        .unique()
    )
    amputation = (
        core__condition.filter(nw.col("normalized_code_type") == "icd-10-cm")
        .filter(~nw.col("encounter_id").is_null())
        .join(
            amp_codes,
            left_on=["normalized_code"],
            right_on=["code"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .unique()
        .with_columns(nw.lit("amputation").alias("exclusion_reason"))
    )

    preg_codes = (
        vs.filter(
            (nw.col("value_set_name") == "mdc_14_principal_diagnosis_codes")
            & (nw.col("pqi_number") == "appendix_f")
        )
        .select("code")
        .unique()
    )
    pregnancy = (
        core__condition.filter(nw.col("normalized_code_type") == "icd-10-cm")
        .filter(~nw.col("encounter_id").is_null())
        .join(
            preg_codes,
            left_on=["normalized_code"],
            right_on=["code"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .unique()
        .with_columns(nw.lit("pregnancy").alias("exclusion_reason"))
    )

    combined = nw.concat(
        [
            ahrq_measures___int_pqi_shared_exclusion_union.select(
                "encounter_id", "data_source", "exclusion_reason"
            ),
            amputation,
            pregnancy,
        ]
    )
    return combined.with_columns(
        nw.col("exclusion_reason")
        .rank(method="ordinal")
        .over("encounter_id", "data_source")
        .cast(nw.Int64)
        .alias("exclusion_number"),
    ).select(
        "encounter_id",
        "data_source",
        "exclusion_reason",
        "exclusion_number",
    )


@nw.narwhalify
def int_pqi_16_num(
    ahrq_measures___stg_pqi_inpatient_encounter: FrameT,
    ahrq_measures___value_set_pqi: FrameT,
    ahrq_measures___int_pqi_16_denom: FrameT,
    ahrq_measures___int_pqi_16_exclusions: FrameT,
    core__condition: FrameT,
    core__procedure: FrameT,
) -> FrameT:
    """PQI 16: diabetes dx + lower-extremity amputation px."""
    vs = ahrq_measures___value_set_pqi
    e = ahrq_measures___stg_pqi_inpatient_encounter

    # Encounters with diabetes diagnosis
    dx_codes = (
        vs.filter(
            (nw.col("value_set_name") == "diabetes_diagnosis_codes")
            & (nw.col("pqi_number") == "16")
        )
        .select("code")
        .unique()
    )
    diagnosis = (
        core__condition.filter(nw.col("normalized_code_type") == "icd-10-cm")
        .filter(~nw.col("encounter_id").is_null())
        .join(
            dx_codes,
            left_on=["normalized_code"],
            right_on=["code"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .unique()
    )

    # Encounters with amputation procedure + diabetes dx
    px_codes = (
        vs.filter(
            (nw.col("value_set_name") == "lower-extremity_amputation_procedure_codes")
            & (nw.col("pqi_number") == "16")
        )
        .select("code")
        .unique()
    )
    procedures = (
        core__procedure.filter(nw.col("normalized_code_type") == "icd-10-pcs")
        .join(
            diagnosis,
            on=["encounter_id", "data_source"],
            how="inner",
        )
        .join(
            px_codes,
            left_on=["normalized_code"],
            right_on=["code"],
            how="inner",
        )
        .select("encounter_id", "data_source")
        .unique()
    )

    result = e.join(
        ahrq_measures___int_pqi_16_denom.select(
            "person_id", "data_source", "year_number"
        ),
        on=["person_id", "data_source", "year_number"],
        how="inner",
    ).join(
        procedures,
        on=["encounter_id", "data_source"],
        how="inner",
    )
    excl = (
        ahrq_measures___int_pqi_16_exclusions.select("encounter_id", "data_source")
        .unique()
        .with_columns(nw.lit(1).alias("_excl"))
    )
    result = result.join(excl, on=["encounter_id", "data_source"], how="left")
    return result.filter(nw.col("_excl").is_null()).select(
        "data_source", "person_id", "year_number", "encounter_id"
    )