stack/notebooks/cms_quality_measures.py

import marimo

__generated_with = "0.21.1"
app = marimo.App(width="full")


@app.cell(hide_code=True)
def _():
    import marimo as mo

    return (mo,)


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    # CMS Quality Measures — Pipeline Explorer

    Interactive walkthrough of three CMS claims-based hospital admission and readmission
    quality measures implemented as **narwhals expression functions**.

    | Measure | Full Name | NQF | Program |
    |---------|-----------|-----|---------|
    | **UAMCC** | All-Cause Unplanned Admissions for Multiple Chronic Conditions | #2888 | ACO REACH + MIPS |
    | **ACR** | Risk-Standardized, All-Condition Readmission | #1789 | ACO REACH |
    | **HWR** | Hospital-wide, 30-Day, All-cause Unplanned Readmission | — | MIPS Groups |

    **Contents:**

    1. Step-by-step walkthrough of each measure's pipeline logic (via `inspect`)
    2. Value set tables loaded from CMS Excel workbooks (PY2024 + PY2025 + PY2026)
    3. Year-over-year diffs between performance years
    4. Pipeline execution on actual Synthea/Tuva claims data
    """)
    return


@app.cell(hide_code=True)
def _():
    import inspect
    import textwrap

    import polars as pl

    from conf import connect
    from conf.display import plain_years

    con = connect.duckdb()

    def q(sql):
        """Run SQL query and return a Polars DataFrame."""
        return con.execute(sql).pl()

    return con, inspect, pl, plain_years, q, textwrap


@app.cell(hide_code=True)
def _():
    from aco.express import cms_quality_measures as ex

    return (ex,)


@app.cell(hide_code=True)
def _(inspect, mo, textwrap):
    def _param_to_table(p):
        """Convert a function parameter name to a qualified table reference."""
        if "___" in p:
            schema, table = p.split("___", 1)
            return f"{schema}._{table}"
        if "__" in p:
            schema, table = p.split("__", 1)
            return f"{schema}.{table}"
        return p

    def fn_doc(fn):
        """Build a marimo markdown block documenting a pipeline function."""
        sig = inspect.signature(fn)
        doc = inspect.getdoc(fn) or "(no docstring)"

        # Parameters → table references
        params = []
        for p in sig.parameters:
            ref = _param_to_table(p)
            params.append(f"| `{p}` | `{ref}` |")
        param_md = "\n".join(params) if params else "| — | — |"

        # Source code (prefer unwrapped original)
        try:
            target = getattr(fn, "__wrapped__", fn)
            src = textwrap.dedent(inspect.getsource(target))
        except (OSError, TypeError):
            src = "(source unavailable)"

        return mo.md(f"""
    **`{fn.__name__}`** `{sig}`

    {doc}

    **Parameters:**

    | Parameter | Table Reference |
    |-----------|----------------|
    {param_md}

    <details><summary>Source code</summary>

    ```python
    {src}```

    </details>
    """)

    return (fn_doc,)


@app.cell(hide_code=True)
def _(ex, inspect, mo):
    _module_doc = inspect.getdoc(ex) or ""
    mo.md(f"""
    ## Module: `aco.express.cms_quality_measures`

    ```
    {_module_doc}
    ```
    """)
    return


@app.cell(hide_code=True)
def _(ex, fn_doc, mo):
    _uamcc_steps = [
        ("0a. Stage Medical Claims", ex.stg_medical_claim),
        ("0b. Stage Claim–Condition Pairs", ex.stg_medical_claim_condition),
        ("1. Performance Period", ex.uamcc_performance_period),
        ("2. MCC Cohort Identification", ex.uamcc_int_mcc_cohort),
        ("3. Denominator", ex.uamcc_int_denominator),
        ("4. Denominator Exclusions", ex.uamcc_int_denominator_exclusion),
        ("5. Planned Admission (PAA v4.0)", ex.uamcc_int_planned_admission),
        ("6. Outcome Exclusions", ex.uamcc_int_outcome_exclusion),
        ("7. Person-Time", ex.uamcc_int_person_time),
        ("8. Numerator", ex.uamcc_int_numerator),
        ("9. Summary", ex.uamcc_summary),
    ]

    mo.vstack(
        [
            mo.md("""## UAMCC Pipeline — 11 Steps
    *All-Cause Unplanned Admissions for Patients with Multiple Chronic Conditions (NQF #2888)*

    MCC-eligible beneficiaries (≥66, 2+ chronic condition groups) →
    exclude hospice / enrollment gaps → classify planned admissions (PAA v4.0) →
    exclude planned / injury / complication admissions → calculate person-time →
    count unplanned admissions → rate per 100 person-years.
    """),
            mo.accordion({title: fn_doc(fn) for title, fn in _uamcc_steps}),
        ]
    )
    return


@app.cell(hide_code=True)
def _(ex, fn_doc, mo):
    _acr_steps = [
        ("1. Performance Period", ex.acr_performance_period),
        ("2. Index Admission", ex.acr_int_index_admission),
        ("3. Specialty Cohort", ex.acr_int_specialty_cohort),
        ("4. Planned Readmission (PAA v4.0)", ex.acr_int_planned_readmission),
        ("5. Summary", ex.acr_summary),
    ]

    mo.vstack(
        [
            mo.md("""## ACR Pipeline — 5 Steps
    *Risk-Standardized, All-Condition Readmission (NQF #1789)*

    Eligible index hospitalizations (≥65, discharged alive, non-transfer) →
    assign to 5 specialty cohorts (Surgery/Gyn, Cardiorespiratory, Cardiovascular,
    Neurology, Medicine) → classify 30-day readmissions as planned / unplanned
    (PAA v4.0) → observed readmission rate.
    """),
            mo.accordion({title: fn_doc(fn) for title, fn in _acr_steps}),
        ]
    )
    return


@app.cell(hide_code=True)
def _(ex, fn_doc, mo):
    _hwr_steps = [
        ("1. Performance Period", ex.hwr_performance_period),
        ("2. Denominator", ex.hwr_int_denominator),
        ("3. Planned Readmission (PAA v4.0)", ex.hwr_int_planned_readmission),
        ("4. Summary", ex.hwr_summary),
    ]

    mo.vstack(
        [
            mo.md("""## HWR Pipeline — 4 Steps
    *Hospital-wide, 30-Day, All-cause Unplanned Readmission (MIPS Groups)*

    Eligible index hospitalizations + AHRQ CCS specialty cohort assignment →
    classify 30-day readmissions (PAA v4.0, same algorithm as ACR) →
    observed readmission rate by TIN clinician group.
    """),
            mo.accordion({title: fn_doc(fn) for title, fn in _hwr_steps}),
        ]
    )
    return


@app.cell(hide_code=True)
def _(con, mo, pl):
    _vs_tables = con.execute("""
        SELECT table_name
        FROM information_schema.tables
        WHERE table_schema = 'cms_quality_measures'
          AND table_name LIKE '%value_set%'
        ORDER BY table_name
    """).fetchall()

    _rows = []
    for (tbl,) in _vs_tables:
        try:
            year_counts = con.execute(f"""
                SELECT performance_year, count(*) AS rows
                FROM cms_quality_measures."{tbl}"
                GROUP BY performance_year
                ORDER BY performance_year
            """).fetchall()
            for yr, cnt in year_counts:
                _rows.append({"table": tbl, "performance_year": yr, "rows": cnt})
        except Exception:
            _rows.append({"table": tbl, "performance_year": None, "rows": 0})

    vs_overview = pl.DataFrame(_rows)

    _vs_pivot = vs_overview.pivot(
        on="performance_year",
        index="table",
        values="rows",
    ).fill_null(0)

    mo.vstack(
        [
            mo.md("## Value Sets — Row Counts by Performance Year"),
            mo.ui.table(_vs_pivot),
        ]
    )
    return


@app.cell(hide_code=True)
def _(mo, q):
    _cohort = q("""
        SELECT performance_year, chronic_condition_group, count(*) AS codes
        FROM cms_quality_measures._uamcc_value_set_cohort
        GROUP BY performance_year, chronic_condition_group
        ORDER BY chronic_condition_group, performance_year
    """)
    _cohort_pivot = _cohort.pivot(
        on="performance_year",
        index="chronic_condition_group",
        values="codes",
    ).fill_null(0)

    _excl = q("""
        SELECT performance_year, exclusion_category, count(*) AS codes
        FROM cms_quality_measures._uamcc_value_set_exclusions
        GROUP BY performance_year, exclusion_category
        ORDER BY exclusion_category, performance_year
    """)
    _excl_pivot = _excl.pivot(
        on="performance_year",
        index="exclusion_category",
        values="codes",
    ).fill_null(0)

    _paa = q("""
        SELECT 'PAA1 — Always Planned Px' AS rule, performance_year, count(*) AS codes
        FROM cms_quality_measures._uamcc_value_set_paa1 GROUP BY performance_year
        UNION ALL
        SELECT 'PAA2 — Always Planned Dx', performance_year, count(*)
        FROM cms_quality_measures._uamcc_value_set_paa2 GROUP BY performance_year
        UNION ALL
        SELECT 'PAA3 — Pot. Planned Px', performance_year, count(*)
        FROM cms_quality_measures._uamcc_value_set_paa3 GROUP BY performance_year
        UNION ALL
        SELECT 'PAA4 — Acute Dx', performance_year, count(*)
        FROM cms_quality_measures._uamcc_value_set_paa4 GROUP BY performance_year
        UNION ALL
        SELECT 'CCS-ICD10CM Crosswalk', performance_year, count(*)
        FROM cms_quality_measures._uamcc_value_set_ccs_icd10_cm GROUP BY performance_year
        UNION ALL
        SELECT 'CCS-ICD10PCS Crosswalk', performance_year, count(*)
        FROM cms_quality_measures._uamcc_value_set_ccs_icd10_pcs GROUP BY performance_year
        ORDER BY rule, performance_year
    """)
    _paa_pivot = _paa.pivot(
        on="performance_year",
        index="rule",
        values="codes",
    ).fill_null(0)

    mo.vstack(
        [
            mo.md("### UAMCC Value Sets"),
            mo.md("**Cohort — ICD-10-CM codes by chronic condition group:**"),
            mo.ui.table(_cohort_pivot),
            mo.md("**Outcome exclusions by category:**"),
            mo.ui.table(_excl_pivot),
            mo.md("**PAA v4.0 rules and CCS crosswalks:**"),
            mo.ui.table(_paa_pivot),
        ]
    )
    return


@app.cell(hide_code=True)
def _(mo, q):
    _acr_ccs = q("""
        SELECT performance_year, specialty_cohort, count(*) AS ccs_categories
        FROM cms_quality_measures._acr_value_set_cohort_ccs
        GROUP BY performance_year, specialty_cohort
        ORDER BY specialty_cohort, performance_year
    """)
    _acr_ccs_pivot = _acr_ccs.pivot(
        on="performance_year",
        index="specialty_cohort",
        values="ccs_categories",
    ).fill_null(0)

    _acr_excl = q("""
        SELECT performance_year, count(*) AS ccs_categories
        FROM cms_quality_measures._acr_value_set_exclusions
        GROUP BY performance_year
        ORDER BY performance_year
    """)

    _acr_paa = q("""
        SELECT 'PAA1 — Always Planned Px' AS rule, performance_year, count(*) AS codes
        FROM cms_quality_measures._acr_value_set_paa1 GROUP BY performance_year
        UNION ALL
        SELECT 'PAA2 — Always Planned Dx', performance_year, count(*)
        FROM cms_quality_measures._acr_value_set_paa2 GROUP BY performance_year
        UNION ALL
        SELECT 'PAA3 — Pot. Planned Px', performance_year, count(*)
        FROM cms_quality_measures._acr_value_set_paa3 GROUP BY performance_year
        UNION ALL
        SELECT 'PAA4 — Acute Dx', performance_year, count(*)
        FROM cms_quality_measures._acr_value_set_paa4 GROUP BY performance_year
        ORDER BY rule, performance_year
    """)
    _acr_paa_pivot = _acr_paa.pivot(
        on="performance_year",
        index="rule",
        values="codes",
    ).fill_null(0)

    _acr_cohort_icd10 = q("""
        SELECT performance_year, count(*) AS icd10_pcs_codes
        FROM cms_quality_measures._acr_value_set_cohort_icd10
        GROUP BY performance_year
        ORDER BY performance_year
    """)

    mo.vstack(
        [
            mo.md("### ACR Value Sets"),
            mo.md("**Cohort CCS categories by specialty cohort:**"),
            mo.ui.table(_acr_ccs_pivot),
            mo.md("**Cohort ICD-10-PCS codes (Surgery/Gyn):**"),
            mo.ui.table(_acr_cohort_icd10),
            mo.md("**Exclusions (CCS diagnosis categories):**"),
            mo.ui.table(_acr_excl),
            mo.md("**PAA v4.0 rules:**"),
            mo.ui.table(_acr_paa_pivot),
        ]
    )
    return


@app.cell(hide_code=True)
def _(mo, q):
    _hwr_spec = q("""
        SELECT performance_year, specialty_cohort, count(*) AS ccs_categories
        FROM cms_quality_measures._hwr_value_set_specialty_cohort
        GROUP BY performance_year, specialty_cohort
        ORDER BY specialty_cohort, performance_year
    """)
    _hwr_spec_pivot = _hwr_spec.pivot(
        on="performance_year",
        index="specialty_cohort",
        values="ccs_categories",
    ).fill_null(0)

    _hwr_excl = q("""
        SELECT performance_year, count(*) AS ccs_categories
        FROM cms_quality_measures._hwr_value_set_cohort_exclusions
        GROUP BY performance_year
        ORDER BY performance_year
    """)

    _hwr_surg = q("""
        SELECT performance_year, count(*) AS icd10_pcs_codes
        FROM cms_quality_measures._hwr_value_set_surg_gyn_cohort
        GROUP BY performance_year
        ORDER BY performance_year
    """)

    _hwr_paa = q("""
        SELECT 'PAA1 — Always Planned Px' AS rule, performance_year, count(*) AS codes
        FROM cms_quality_measures._hwr_value_set_paa1 GROUP BY performance_year
        UNION ALL
        SELECT 'PAA2 — Always Planned Dx', performance_year, count(*)
        FROM cms_quality_measures._hwr_value_set_paa2 GROUP BY performance_year
        UNION ALL
        SELECT 'PAA3 — Pot. Planned Px', performance_year, count(*)
        FROM cms_quality_measures._hwr_value_set_paa3 GROUP BY performance_year
        UNION ALL
        SELECT 'PAA4 — Acute Dx', performance_year, count(*)
        FROM cms_quality_measures._hwr_value_set_paa4 GROUP BY performance_year
        ORDER BY rule, performance_year
    """)
    _hwr_paa_pivot = _hwr_paa.pivot(
        on="performance_year",
        index="rule",
        values="codes",
    ).fill_null(0)

    mo.vstack(
        [
            mo.md("### HWR Value Sets"),
            mo.md("**Specialty cohort CCS categories:**"),
            mo.ui.table(_hwr_spec_pivot),
            mo.md("**Surgery/Gynecology ICD-10-PCS codes:**"),
            mo.ui.table(_hwr_surg),
            mo.md("**Cohort exclusions:**"),
            mo.ui.table(_hwr_excl),
            mo.md("**PAA v4.0 rules:**"),
            mo.ui.table(_hwr_paa_pivot),
        ]
    )
    return


@app.cell(hide_code=True)
def _(con, mo, pl):
    _DIFF_SPECS = [
        ("_uamcc_value_set_cohort", "icd_10_cm"),
        ("_uamcc_value_set_exclusions", "category_or_code"),
        ("_uamcc_value_set_paa1", "ccs_procedure_category"),
        ("_uamcc_value_set_paa2", "ccs_diagnosis_category"),
        ("_uamcc_value_set_paa3", "category_or_code"),
        ("_uamcc_value_set_paa4", "category_or_code"),
        ("_uamcc_value_set_ccs_icd10_cm", "icd_10_cm"),
        ("_uamcc_value_set_ccs_icd10_pcs", "icd_10_pcs"),
        ("_acr_value_set_cohort_ccs", "ccs_category"),
        ("_acr_value_set_cohort_icd10", "icd_10_pcs"),
        ("_acr_value_set_exclusions", "ccs_diagnosis_category"),
        ("_acr_value_set_paa1", "ccs_procedure_category"),
        ("_acr_value_set_paa2", "ccs_diagnosis_category"),
        ("_acr_value_set_paa3", "category_or_code"),
        ("_acr_value_set_paa4", "category_or_code"),
        ("_hwr_value_set_specialty_cohort", "ccs_category"),
        ("_hwr_value_set_surg_gyn_cohort", "icd_10_pcs"),
        ("_hwr_value_set_cohort_exclusions", "ccs_diagnosis_category"),
        ("_hwr_value_set_paa1", "ccs_procedure_category"),
        ("_hwr_value_set_paa2", "ccs_diagnosis_category"),
        ("_hwr_value_set_paa3", "category_or_code"),
        ("_hwr_value_set_paa4", "category_or_code"),
    ]

    def _normalize_code(val):
        """Strip dots and lowercase so PY2024 'I21.01' matches PY2025 'I2101'."""
        if val is None:
            return None
        return val.replace(".", "").upper()

    _diff_rows = []
    for _tbl, _key_col in _DIFF_SPECS:
        _qualified = f'cms_quality_measures."{_tbl}"'
        try:
            _years = [
                r[0]
                for r in con.execute(
                    f"SELECT DISTINCT performance_year FROM {_qualified} ORDER BY 1"
                ).fetchall()
            ]
        except Exception:
            continue
        for _i in range(len(_years) - 1):
            _ya, _yb = _years[_i], _years[_i + 1]
            _set_a = {
                _normalize_code(r[0])
                for r in con.execute(
                    f'SELECT DISTINCT "{_key_col}" FROM {_qualified}'
                    f" WHERE performance_year = {_ya}"
                    f' AND "{_key_col}" IS NOT NULL'
                ).fetchall()
            } - {None}
            _set_b = {
                _normalize_code(r[0])
                for r in con.execute(
                    f'SELECT DISTINCT "{_key_col}" FROM {_qualified}'
                    f" WHERE performance_year = {_yb}"
                    f' AND "{_key_col}" IS NOT NULL'
                ).fetchall()
            } - {None}
            _added = len(_set_b - _set_a)
            _removed = len(_set_a - _set_b)
            _unchanged = len(_set_a & _set_b)
            _total = len(_set_a) + _added
            _pct = round((_added + _removed) / _total * 100, 1) if _total else 0
            _diff_rows.append(
                {
                    "table": _tbl,
                    "key_column": _key_col,
                    "years": f"{_ya} -> {_yb}",
                    "in_a": len(_set_a),
                    "in_b": len(_set_b),
                    "added": _added,
                    "removed": _removed,
                    "unchanged": _unchanged,
                    "pct_changed": _pct,
                }
            )

    diff_df = pl.DataFrame(_diff_rows)

    _total_added = diff_df["added"].sum()
    _total_removed = diff_df["removed"].sum()
    _tables_changed = diff_df.filter(
        (pl.col("added") > 0) | (pl.col("removed") > 0)
    ).height

    mo.vstack(
        [
            mo.md(f"""## Year-over-Year Value Set Diffs

    **{_tables_changed}** of {diff_df.height} value set comparisons show changes:
    **+{_total_added}** codes added, **-{_total_removed}** removed.
    """),
            mo.ui.table(diff_df),
        ]
    )
    return (diff_df,)


@app.cell(hide_code=True)
def _(con, diff_df, mo):
    def _normalize_code(val):
        if val is None:
            return None
        return val.replace(".", "").upper()

    _detail_items = {}
    for _row in diff_df.iter_rows(named=True):
        if _row["added"] == 0 and _row["removed"] == 0:
            continue
        _tbl = _row["table"]
        _key_col = _row["key_column"]
        _qualified = f'cms_quality_measures."{_tbl}"'

        # Parse years from "2025 -> 2026" format
        _ya_str, _yb_str = _row["years"].split(" -> ")
        _ya, _yb = int(_ya_str), int(_yb_str)

        # Build normalized→original mappings for display
        _raw_a = {
            r[0]
            for r in con.execute(
                f'SELECT DISTINCT "{_key_col}" FROM {_qualified}'
                f" WHERE performance_year = {_ya}"
                f' AND "{_key_col}" IS NOT NULL'
            ).fetchall()
        }
        _raw_b = {
            r[0]
            for r in con.execute(
                f'SELECT DISTINCT "{_key_col}" FROM {_qualified}'
                f" WHERE performance_year = {_yb}"
                f' AND "{_key_col}" IS NOT NULL'
            ).fetchall()
        }
        _set_a = {_normalize_code(v) for v in _raw_a}
        _set_b = {_normalize_code(v) for v in _raw_b}
        # Map normalized back to original (prefer the newer year's form)
        _norm_to_orig = {_normalize_code(v): v for v in _raw_a}
        _norm_to_orig.update({_normalize_code(v): v for v in _raw_b})

        _added_codes = sorted(
            _norm_to_orig[c] for c in (_set_b - _set_a) if c is not None
        )
        _removed_codes = sorted(
            _norm_to_orig[c] for c in (_set_a - _set_b) if c is not None
        )
        _parts = []
        if _added_codes:
            _sample = _added_codes[:20]
            _parts.append(
                f"**Added** ({len(_added_codes)}): "
                + ", ".join(f"`{c}`" for c in _sample)
                + ("..." if len(_added_codes) > 20 else "")
            )
        if _removed_codes:
            _sample = _removed_codes[:20]
            _parts.append(
                f"**Removed** ({len(_removed_codes)}): "
                + ", ".join(f"`{c}`" for c in _sample)
                + ("..." if len(_removed_codes) > 20 else "")
            )
        _detail_items[f"{_tbl} ({_row['years']})"] = mo.md("\n\n".join(_parts))

    mo.vstack(
        [
            mo.md("### Diff Detail — Codes Added / Removed"),
            mo.accordion(_detail_items)
            if _detail_items
            else mo.md("No changes detected."),
        ]
    )
    return


@app.cell(hide_code=True)
def _(con, inspect, mo, pl):
    from datetime import date

    # CY2018 — latest full year in the Synthea/Tuva data
    PERF_YEAR = 2018
    PERF_BEGIN = date(2018, 1, 1)
    PERF_END = date(2018, 12, 31)
    LOOKBACK_BEGIN = date(2017, 1, 1)
    LOOKBACK_END = date(2017, 12, 31)

    # Performance period anchor DataFrames
    uamcc_pp = pl.DataFrame(
        {
            "measure_id": ["UAMCC"],
            "measure_name": ["All-Cause Unplanned Admissions for MCC"],
            "nqf_id": ["2888"],
            "performance_year": [PERF_YEAR],
            "performance_period_begin": [PERF_BEGIN],
            "performance_period_end": [PERF_END],
            "lookback_period_begin": [LOOKBACK_BEGIN],
            "lookback_period_end": [LOOKBACK_END],
        }
    )

    acr_pp = pl.DataFrame(
        {
            "measure_id": ["ACR"],
            "measure_name": ["Risk-Standardized All-Condition Readmission"],
            "nqf_id": ["1789"],
            "performance_year": [PERF_YEAR],
            "performance_period_begin": [PERF_BEGIN],
            "performance_period_end": [PERF_END],
        }
    )

    hwr_pp = pl.DataFrame(
        {
            "measure_id": ["HWR"],
            "measure_name": ["Hospital-Wide 30-Day All-Cause Unplanned Readmission"],
            "performance_year": [PERF_YEAR],
            "performance_period_begin": [PERF_BEGIN],
            "performance_period_end": [PERF_END],
        }
    )

    # Table ref fallbacks for missing ACR CCS crosswalks (reuse UAMCC versions)
    _TABLE_FALLBACKS = {
        "cms_quality_measures._acr_value_set_ccs_icd10_cm": "cms_quality_measures._uamcc_value_set_ccs_icd10_cm",
        "cms_quality_measures._acr_value_set_ccs_icd10_pcs": "cms_quality_measures._uamcc_value_set_ccs_icd10_pcs",
    }

    _load_cache: dict[str, pl.DataFrame] = {}

    def load(table_ref):
        """Load a table from DuckDB, with caching and fallbacks."""
        if table_ref in _load_cache:
            return _load_cache[table_ref]
        _schema, _tbl = table_ref.split(".", 1)
        try:
            _df = con.execute(f'SELECT * FROM "{_schema}"."{_tbl}"').pl()
        except Exception:
            _fallback = _TABLE_FALLBACKS.get(table_ref)
            if _fallback:
                _fs, _ft = _fallback.split(".", 1)
                _df = con.execute(f'SELECT * FROM "{_fs}"."{_ft}"').pl()
            else:
                raise
        # For value set tables, use latest year to avoid join duplicates
        if "value_set" in _tbl and "performance_year" in _df.columns:
            _max_yr = _df["performance_year"].max()
            _df = _df.filter(pl.col("performance_year") == _max_yr)
        _load_cache[table_ref] = _df
        return _df

    def _param_to_table_ref(p):
        if "___" in p:
            schema, table = p.split("___", 1)
            return f"{schema}._{table}"
        if "__" in p:
            schema, table = p.split("__", 1)
            return f"{schema}.{table}"
        return p

    def run_step(fn, cache):
        """Execute one pipeline function, resolving params from cache or DB."""
        sig = inspect.signature(fn)
        kwargs = {}
        for param in sig.parameters:
            ref = _param_to_table_ref(param)
            if ref in cache:
                kwargs[param] = cache[ref]
            else:
                kwargs[param] = load(ref)
        return fn(**kwargs)

    mo.md(f"""
    ## Pipeline Execution on Actual Data

    Running all three measure pipelines against **Synthea/Tuva claims data**
    (1,000 patients, CY{PERF_YEAR} measurement period).

    - Performance period: **{PERF_BEGIN}** to **{PERF_END}**
    - Lookback window: **{LOOKBACK_BEGIN}** to **{LOOKBACK_END}**
    """)
    return PERF_YEAR, acr_pp, hwr_pp, run_step, uamcc_pp


@app.cell(hide_code=True)
def _(ex, mo, pl, plain_years, run_step, uamcc_pp):
    _uamcc_pipeline = [
        # Staging tables (_stg_medical_claim, _stg_medical_claim_condition) are
        # materialized in DuckDB and loaded on demand by load().
        ("cms_quality_measures._uamcc_performance_period", ex.uamcc_performance_period),
        ("cms_quality_measures._uamcc_int_mcc_cohort", ex.uamcc_int_mcc_cohort),
        ("cms_quality_measures._uamcc_int_denominator", ex.uamcc_int_denominator),
        (
            "cms_quality_measures._uamcc_int_denominator_exclusion",
            ex.uamcc_int_denominator_exclusion,
        ),
        (
            "cms_quality_measures._uamcc_int_planned_admission",
            ex.uamcc_int_planned_admission,
        ),
        (
            "cms_quality_measures._uamcc_int_outcome_exclusion",
            ex.uamcc_int_outcome_exclusion,
        ),
        ("cms_quality_measures._uamcc_int_person_time", ex.uamcc_int_person_time),
        ("cms_quality_measures._uamcc_int_numerator", ex.uamcc_int_numerator),
        ("cms_quality_measures.uamcc_summary", ex.uamcc_summary),
    ]

    _uamcc_cache = {
        "cms_quality_measures._uamcc_performance_period": uamcc_pp,
    }
    uamcc_results = {}
    _uamcc_errors = {}

    for _name, _fn in _uamcc_pipeline:
        try:
            _result = run_step(_fn, _uamcc_cache)
            _uamcc_cache[_name] = _result
            uamcc_results[_name] = _result
        except Exception as _exc:
            _uamcc_errors[_name] = str(_exc)

    _step_info = []
    for _name, _ in _uamcc_pipeline:
        _short = _name.split(".")[-1]
        if _name in uamcc_results:
            _df = uamcc_results[_name]
            _step_info.append(
                {"step": _short, "rows": len(_df), "columns": len(_df.columns)}
            )
        elif _name in _uamcc_errors:
            _step_info.append({"step": _short, "rows": -1, "columns": 0})
    _uamcc_step_df = pl.DataFrame(_step_info)

    _items = {}
    for _name, _ in _uamcc_pipeline:
        _short = _name.split(".")[-1]
        if _name in uamcc_results:
            _df = uamcc_results[_name]
            _items[f"{_short} ({len(_df)} rows)"] = mo.ui.table(plain_years(_df.head(50)))
        elif _name in _uamcc_errors:
            _items[f"{_short} (ERROR)"] = mo.md(f"```\n{_uamcc_errors[_name]}\n```")

    mo.vstack(
        [
            mo.md("### UAMCC Pipeline Results"),
            mo.ui.table(_uamcc_step_df, label="Step Summary"),
            mo.accordion(_items),
        ]
    )
    return (uamcc_results,)


@app.cell(hide_code=True)
def _(acr_pp, ex, mo, pl, plain_years, run_step):
    _acr_pipeline = [
        ("cms_quality_measures._acr_performance_period", ex.acr_performance_period),
        ("cms_quality_measures._acr_int_index_admission", ex.acr_int_index_admission),
        ("cms_quality_measures._acr_int_specialty_cohort", ex.acr_int_specialty_cohort),
        (
            "cms_quality_measures._acr_int_planned_readmission",
            ex.acr_int_planned_readmission,
        ),
        ("cms_quality_measures.acr_summary", ex.acr_summary),
    ]

    _acr_cache = {
        "cms_quality_measures._acr_performance_period": acr_pp,
    }
    acr_results = {}
    _acr_errors = {}

    for _name, _fn in _acr_pipeline:
        try:
            _result = run_step(_fn, _acr_cache)
            _acr_cache[_name] = _result
            acr_results[_name] = _result
        except Exception as _exc:
            _acr_errors[_name] = str(_exc)

    _step_info = []
    for _name, _ in _acr_pipeline:
        _short = _name.split(".")[-1]
        if _name in acr_results:
            _df = acr_results[_name]
            _step_info.append(
                {"step": _short, "rows": len(_df), "columns": len(_df.columns)}
            )
        elif _name in _acr_errors:
            _step_info.append({"step": _short, "rows": -1, "columns": 0})
    _acr_step_df = pl.DataFrame(_step_info)

    _items = {}
    for _name, _ in _acr_pipeline:
        _short = _name.split(".")[-1]
        if _name in acr_results:
            _df = acr_results[_name]
            _items[f"{_short} ({len(_df)} rows)"] = mo.ui.table(plain_years(_df.head(50)))
        elif _name in _acr_errors:
            _items[f"{_short} (ERROR)"] = mo.md(f"```\n{_acr_errors[_name]}\n```")

    mo.vstack(
        [
            mo.md("### ACR Pipeline Results"),
            mo.ui.table(_acr_step_df, label="Step Summary"),
            mo.accordion(_items),
        ]
    )
    return (acr_results,)


@app.cell(hide_code=True)
def _(ex, hwr_pp, mo, pl, plain_years, run_step):
    _hwr_pipeline = [
        ("cms_quality_measures._hwr_performance_period", ex.hwr_performance_period),
        ("cms_quality_measures._hwr_int_denominator", ex.hwr_int_denominator),
        (
            "cms_quality_measures._hwr_int_planned_readmission",
            ex.hwr_int_planned_readmission,
        ),
        ("cms_quality_measures.hwr_summary", ex.hwr_summary),
    ]

    _hwr_cache = {
        "cms_quality_measures._hwr_performance_period": hwr_pp,
    }
    hwr_results = {}
    _hwr_errors = {}

    for _name, _fn in _hwr_pipeline:
        try:
            _result = run_step(_fn, _hwr_cache)
            _hwr_cache[_name] = _result
            hwr_results[_name] = _result
        except Exception as _exc:
            _hwr_errors[_name] = str(_exc)

    _step_info = []
    for _name, _ in _hwr_pipeline:
        _short = _name.split(".")[-1]
        if _name in hwr_results:
            _df = hwr_results[_name]
            _step_info.append(
                {"step": _short, "rows": len(_df), "columns": len(_df.columns)}
            )
        elif _name in _hwr_errors:
            _step_info.append({"step": _short, "rows": -1, "columns": 0})
    _hwr_step_df = pl.DataFrame(_step_info)

    _items = {}
    for _name, _ in _hwr_pipeline:
        _short = _name.split(".")[-1]
        if _name in hwr_results:
            _df = hwr_results[_name]
            _items[f"{_short} ({len(_df)} rows)"] = mo.ui.table(plain_years(_df.head(50)))
        elif _name in _hwr_errors:
            _items[f"{_short} (ERROR)"] = mo.md(f"```\n{_hwr_errors[_name]}\n```")

    mo.vstack(
        [
            mo.md("### HWR Pipeline Results"),
            mo.ui.table(_hwr_step_df, label="Step Summary"),
            mo.accordion(_items),
        ]
    )
    return (hwr_results,)


@app.cell(hide_code=True)
def _(PERF_YEAR, acr_results, hwr_results, mo, pl, uamcc_results):
    def _safe(df, col):
        try:
            v = df[col][0]
            return v if v is not None else None
        except Exception:
            return None

    _uamcc_s = uamcc_results.get("cms_quality_measures.uamcc_summary")
    _acr_s = acr_results.get("cms_quality_measures.acr_summary")
    _hwr_s = hwr_results.get("cms_quality_measures.hwr_summary")

    _summary_rows = []
    if _uamcc_s is not None:
        _summary_rows.append(
            {
                "measure": "UAMCC",
                "performance_year": PERF_YEAR,
                "denominator": _safe(_uamcc_s, "denominator_count"),
                "numerator": _safe(_uamcc_s, "observed_admissions"),
                "person_years": _safe(_uamcc_s, "total_person_years"),
                "observed_rate": _safe(_uamcc_s, "observed_rate_per_100"),
                "unit": "per 100 person-years",
            }
        )
    if _acr_s is not None:
        _summary_rows.append(
            {
                "measure": "ACR",
                "performance_year": PERF_YEAR,
                "denominator": _safe(_acr_s, "denominator_count"),
                "numerator": _safe(_acr_s, "observed_readmissions"),
                "person_years": None,
                "observed_rate": _safe(_acr_s, "observed_rate"),
                "unit": "proportion",
            }
        )
    if _hwr_s is not None:
        _summary_rows.append(
            {
                "measure": "HWR",
                "performance_year": PERF_YEAR,
                "denominator": _safe(_hwr_s, "denominator_count"),
                "numerator": _safe(_hwr_s, "observed_readmissions"),
                "person_years": None,
                "observed_rate": _safe(_hwr_s, "observed_rate"),
                "unit": "proportion",
            }
        )

    if _summary_rows:
        _summary_df = pl.DataFrame(_summary_rows)
    else:
        _summary_df = pl.DataFrame({"measure": [], "note": ["No results available"]})

    mo.vstack(
        [
            mo.md("""## Cross-Measure Summary

    Observed rates from running all three measure pipelines on Synthea claims data.
    These are **crude observed rates** — the full risk-standardized rates (RSAAR / RSRR)
    require CMS's hierarchical models fit across all ACO data nationally.
    """),
            mo.ui.table(_summary_df),
        ]
    )
    return


@app.cell(hide_code=True)
def _(con, ex, inspect, mo, pl):
    from datetime import date as _date

    # ── Discover available spec years per value-set table ────────────────
    _vs_tables = con.execute("""
        SELECT table_name
        FROM information_schema.tables
        WHERE table_schema = 'cms_quality_measures'
          AND table_name LIKE '%value_set%'
    """).fetchall()

    _table_years: dict[str, set[int]] = {}
    for (_tbl,) in _vs_tables:
        _cols = [
            r[0]
            for r in con.execute(
                f"SELECT column_name FROM information_schema.columns "
                f"WHERE table_schema = 'cms_quality_measures' AND table_name = '{_tbl}'"
            ).fetchall()
        ]
        if "performance_year" in _cols:
            _yrs = con.execute(
                f'SELECT DISTINCT performance_year FROM "cms_quality_measures"."{_tbl}"'
            ).fetchall()
            _table_years[f"cms_quality_measures.{_tbl}"] = {y[0] for y in _yrs}

    # ── Fallbacks: ACR CCS → UAMCC CCS ─────────────────────────────────
    _FALLBACKS = {
        "cms_quality_measures._acr_value_set_ccs_icd10_cm": "cms_quality_measures._uamcc_value_set_ccs_icd10_cm",
        "cms_quality_measures._acr_value_set_ccs_icd10_pcs": "cms_quality_measures._uamcc_value_set_ccs_icd10_pcs",
    }

    # ── Helper: load a table filtered to a specific spec year ───────────
    _year_cache: dict[tuple[str, int], pl.DataFrame] = {}
    _fallback_log: list[dict] = []

    def _load_for_year(table_ref: str, spec_year: int) -> pl.DataFrame:
        _key = (table_ref, spec_year)
        if _key in _year_cache:
            return _year_cache[_key]

        _ref = _FALLBACKS.get(table_ref, table_ref)
        _schema, _tbl = _ref.split(".", 1)
        _df = con.execute(f'SELECT * FROM "{_schema}"."{_tbl}"').pl()

        if "performance_year" in _df.columns:
            avail = sorted(_df["performance_year"].unique().to_list())
            if spec_year in avail:
                _df = _df.filter(pl.col("performance_year") == spec_year)
            else:
                nearest = min(avail, key=lambda y: abs(y - spec_year))
                _df = _df.filter(pl.col("performance_year") == nearest)
                _fallback_log.append(
                    {
                        "table": table_ref.split(".")[-1],
                        "requested_year": spec_year,
                        "actual_year": nearest,
                    }
                )

        _year_cache[_key] = _df
        return _df

    # ── Helper: resolve params and run one step ─────────────────────────
    def _param_to_ref(p):
        if "___" in p:
            s, t = p.split("___", 1)
            return f"{s}._{t}"
        if "__" in p:
            s, t = p.split("__", 1)
            return f"{s}.{t}"
        return p

    def _run_step_for_year(fn, cache, spec_year):
        sig = inspect.signature(fn)
        kwargs = {}
        for param in sig.parameters:
            ref = _param_to_ref(param)
            if ref in cache:
                kwargs[param] = cache[ref]
            else:
                kwargs[param] = _load_for_year(ref, spec_year)
        return fn(**kwargs)

    # ── Define pipelines ────────────────────────────────────────────────
    _PERF_YEAR = 2018
    _PERF_BEGIN = _date(2018, 1, 1)
    _PERF_END = _date(2018, 12, 31)
    _LOOKBACK_BEGIN = _date(2017, 1, 1)
    _LOOKBACK_END = _date(2017, 12, 31)

    _uamcc_pp = pl.DataFrame(
        {
            "measure_id": ["UAMCC"],
            "measure_name": ["All-Cause Unplanned Admissions for MCC"],
            "nqf_id": ["2888"],
            "performance_year": [_PERF_YEAR],
            "performance_period_begin": [_PERF_BEGIN],
            "performance_period_end": [_PERF_END],
            "lookback_period_begin": [_LOOKBACK_BEGIN],
            "lookback_period_end": [_LOOKBACK_END],
        }
    )
    _acr_pp = pl.DataFrame(
        {
            "measure_id": ["ACR"],
            "measure_name": ["Risk-Standardized All-Condition Readmission"],
            "nqf_id": ["1789"],
            "performance_year": [_PERF_YEAR],
            "performance_period_begin": [_PERF_BEGIN],
            "performance_period_end": [_PERF_END],
        }
    )
    _hwr_pp = pl.DataFrame(
        {
            "measure_id": ["HWR"],
            "measure_name": ["Hospital-Wide 30-Day All-Cause Unplanned Readmission"],
            "performance_year": [_PERF_YEAR],
            "performance_period_begin": [_PERF_BEGIN],
            "performance_period_end": [_PERF_END],
        }
    )

    _pipelines = {
        "UAMCC": {
            "steps": [
                (
                    "cms_quality_measures._uamcc_performance_period",
                    ex.uamcc_performance_period,
                ),
                ("cms_quality_measures._uamcc_int_mcc_cohort", ex.uamcc_int_mcc_cohort),
                (
                    "cms_quality_measures._uamcc_int_denominator",
                    ex.uamcc_int_denominator,
                ),
                (
                    "cms_quality_measures._uamcc_int_denominator_exclusion",
                    ex.uamcc_int_denominator_exclusion,
                ),
                (
                    "cms_quality_measures._uamcc_int_planned_admission",
                    ex.uamcc_int_planned_admission,
                ),
                (
                    "cms_quality_measures._uamcc_int_outcome_exclusion",
                    ex.uamcc_int_outcome_exclusion,
                ),
                (
                    "cms_quality_measures._uamcc_int_person_time",
                    ex.uamcc_int_person_time,
                ),
                ("cms_quality_measures._uamcc_int_numerator", ex.uamcc_int_numerator),
                ("cms_quality_measures.uamcc_summary", ex.uamcc_summary),
            ],
            "pp_key": "cms_quality_measures._uamcc_performance_period",
            "pp_df": _uamcc_pp,
            "summary_key": "cms_quality_measures.uamcc_summary",
        },
        "ACR": {
            "steps": [
                (
                    "cms_quality_measures._acr_performance_period",
                    ex.acr_performance_period,
                ),
                (
                    "cms_quality_measures._acr_int_index_admission",
                    ex.acr_int_index_admission,
                ),
                (
                    "cms_quality_measures._acr_int_specialty_cohort",
                    ex.acr_int_specialty_cohort,
                ),
                (
                    "cms_quality_measures._acr_int_planned_readmission",
                    ex.acr_int_planned_readmission,
                ),
                ("cms_quality_measures.acr_summary", ex.acr_summary),
            ],
            "pp_key": "cms_quality_measures._acr_performance_period",
            "pp_df": _acr_pp,
            "summary_key": "cms_quality_measures.acr_summary",
        },
        "HWR": {
            "steps": [
                (
                    "cms_quality_measures._hwr_performance_period",
                    ex.hwr_performance_period,
                ),
                ("cms_quality_measures._hwr_int_denominator", ex.hwr_int_denominator),
                (
                    "cms_quality_measures._hwr_int_planned_readmission",
                    ex.hwr_int_planned_readmission,
                ),
                ("cms_quality_measures.hwr_summary", ex.hwr_summary),
            ],
            "pp_key": "cms_quality_measures._hwr_performance_period",
            "pp_df": _hwr_pp,
            "summary_key": "cms_quality_measures.hwr_summary",
        },
    }

    # ── Determine available spec years per measure ──────────────────────
    # A spec year is available if every value-set table the measure needs
    # has data for that year (or a fallback table does).
    def _measure_vs_tables(measure_cfg):
        """Collect all value_set table refs used by any step in the pipeline."""
        refs = set()
        for _name, fn in measure_cfg["steps"]:
            sig = inspect.signature(fn)
            for p in sig.parameters:
                ref = _param_to_ref(p)
                if "value_set" in ref:
                    resolved = _FALLBACKS.get(ref, ref)
                    refs.add(resolved)
        return refs

    _measure_years: dict[str, list[int]] = {}
    for _mname, _mcfg in _pipelines.items():
        _refs = _measure_vs_tables(_mcfg)
        if _refs:
            _common = set.intersection(*(_table_years.get(r, set()) for r in _refs))
            _measure_years[_mname] = sorted(_common)
        else:
            _measure_years[_mname] = []

    # ── Run each measure × spec year ────────────────────────────────────
    _sensitivity_rows = []
    _step_detail_rows = []
    for _mname, _mcfg in _pipelines.items():
        for _sy in _measure_years[_mname]:
            _cache = {_mcfg["pp_key"]: _mcfg["pp_df"]}
            _error = None
            for _step_name, _fn in _mcfg["steps"]:
                try:
                    _result = _run_step_for_year(_fn, _cache, _sy)
                    _cache[_step_name] = _result
                    _step_detail_rows.append(
                        {
                            "measure": _mname,
                            "spec_year": _sy,
                            "step": _step_name.split(".")[-1],
                            "rows": len(_result),
                        }
                    )
                except Exception as _exc:
                    _error = str(_exc)
                    _step_detail_rows.append(
                        {
                            "measure": _mname,
                            "spec_year": _sy,
                            "step": _step_name.split(".")[-1],
                            "rows": -1,
                        }
                    )
                    break

            _summary = _cache.get(_mcfg["summary_key"])
            if _summary is not None and _error is None:
                _row = {"measure": _mname, "spec_year": _sy}
                for c in _summary.columns:
                    _row[c] = _summary[c][0]
                _sensitivity_rows.append(_row)
            else:
                _sensitivity_rows.append(
                    {
                        "measure": _mname,
                        "spec_year": _sy,
                        "error": _error or "no summary",
                    }
                )

    if _sensitivity_rows:
        _sens_df = pl.DataFrame(_sensitivity_rows)
    else:
        _sens_df = pl.DataFrame({"note": ["No results"]})

    # ── Build a concise comparison view ─────────────────────────────────
    _display_cols = ["measure", "spec_year"]
    _optional = [
        "denominator_count",
        "observed_admissions",
        "total_person_years",
        "observed_rate_per_100",
        "observed_readmissions",
        "observed_rate",
        "error",
    ]
    for _c in _optional:
        if _c in _sens_df.columns:
            _display_cols.append(_c)
    _sens_display = _sens_df.select([c for c in _display_cols if c in _sens_df.columns])

    # Step-level detail pivoted: step × spec_year → rows
    _step_df = pl.DataFrame(_step_detail_rows)
    _step_pivots = {}
    for _m in _step_df["measure"].unique().to_list():
        _msteps = _step_df.filter(pl.col("measure") == _m)
        _piv = _msteps.pivot(on="spec_year", index="step", values="rows")
        _step_pivots[_m] = _piv

    _year_list_md = "\n".join(
        f"- **{m}**: spec years {yrs}" for m, yrs in _measure_years.items()
    )

    # Fallback warnings
    if _fallback_log:
        _fb_lines = [
            f"  - `{fb['table']}`: requested {fb['requested_year']}, used **{fb['actual_year']}**"
            for fb in _fallback_log
        ]
        _fb_md = "\n**Nearest-year fallbacks used:**\n" + "\n".join(_fb_lines) + "\n"
    else:
        _fb_md = ""

    _items = {}
    for _m, _piv in _step_pivots.items():
        _items[f"{_m} — row counts per step"] = mo.ui.table(_piv)

    mo.vstack(
        [
            mo.md(f"""## Spec Year Sensitivity Analysis

    How do year-over-year changes in CMS value set specifications affect measure
    results on the **same population** (Synthea CY{_PERF_YEAR})?

    Each measure is re-run using value sets from each available performance year spec,
    with nearest-year fallback for tables missing a specific year.

    **Available spec years per measure:**

    {_year_list_md}
    {_fb_md}
    > **Note:** Identical results across spec years likely mean that the ~5,600 unique
    > diagnosis codes in the Synthea synthetic population don't overlap with the codes
    > that CMS added or removed between spec years. With real-world claims data covering
    > a broader code space, spec year changes would be more likely to produce observable
    > differences in measure outcomes.
    """),
            mo.ui.table(_sens_display, label="Summary by Spec Year"),
            mo.accordion(_items),
        ]
    )
    return


@app.cell(hide_code=True)
def _(con, mo, pl):
    def _norm(v):
        return v.replace(".", "").upper() if v else None

    # Every value-set table, its code column, measure, and functional role.
    _VS = [
        (
            "_uamcc_value_set_cohort",
            "icd_10_cm",
            "UAMCC",
            "MCC cohort inclusion",
            "ICD-10-CM",
        ),
        (
            "_uamcc_value_set_exclusions",
            "category_or_code",
            "UAMCC",
            "Outcome exclusion",
            "CCS/ICD-10",
        ),
        (
            "_uamcc_value_set_paa1",
            "ccs_procedure_category",
            "UAMCC",
            "PAA Rule 1 — always-planned procedure CCS",
            "CCS",
        ),
        (
            "_uamcc_value_set_paa2",
            "ccs_diagnosis_category",
            "UAMCC",
            "PAA Rule 2 — always-planned diagnosis CCS",
            "CCS",
        ),
        (
            "_uamcc_value_set_paa3",
            "category_or_code",
            "UAMCC",
            "PAA Rule 3 — potentially-planned procedure",
            "CCS/ICD-10-PCS",
        ),
        (
            "_uamcc_value_set_paa4",
            "category_or_code",
            "UAMCC",
            "PAA Rule 3 gate — acute diagnosis",
            "CCS/ICD-10-CM",
        ),
        (
            "_uamcc_value_set_ccs_icd10_cm",
            "icd_10_cm",
            "UAMCC",
            "CCS crosswalk (diagnosis)",
            "ICD-10-CM",
        ),
        (
            "_uamcc_value_set_ccs_icd10_pcs",
            "icd_10_pcs",
            "UAMCC",
            "CCS crosswalk (procedure)",
            "ICD-10-PCS",
        ),
        (
            "_acr_value_set_cohort_ccs",
            "ccs_category",
            "ACR",
            "Specialty cohort CCS",
            "CCS",
        ),
        (
            "_acr_value_set_cohort_icd10",
            "icd_10_pcs",
            "ACR",
            "Specialty cohort ICD-10-PCS",
            "ICD-10-PCS",
        ),
        (
            "_acr_value_set_exclusions",
            "ccs_diagnosis_category",
            "ACR",
            "Cohort exclusion CCS",
            "CCS",
        ),
        (
            "_acr_value_set_paa1",
            "ccs_procedure_category",
            "ACR",
            "PAA Rule 1 — always-planned procedure CCS",
            "CCS",
        ),
        (
            "_acr_value_set_paa2",
            "ccs_diagnosis_category",
            "ACR",
            "PAA Rule 2 — always-planned diagnosis CCS",
            "CCS",
        ),
        (
            "_acr_value_set_paa3",
            "category_or_code",
            "ACR",
            "PAA Rule 3 — potentially-planned procedure",
            "CCS/ICD-10-PCS",
        ),
        (
            "_acr_value_set_paa4",
            "category_or_code",
            "ACR",
            "PAA Rule 3 gate — acute diagnosis",
            "CCS/ICD-10-CM",
        ),
        (
            "_hwr_value_set_specialty_cohort",
            "ccs_category",
            "HWR",
            "Specialty cohort CCS",
            "CCS",
        ),
        (
            "_hwr_value_set_surg_gyn_cohort",
            "icd_10_pcs",
            "HWR",
            "Surgery/Gyn cohort ICD-10-PCS",
            "ICD-10-PCS",
        ),
        (
            "_hwr_value_set_cohort_exclusions",
            "ccs_diagnosis_category",
            "HWR",
            "Cohort exclusion CCS",
            "CCS",
        ),
        (
            "_hwr_value_set_paa1",
            "ccs_procedure_category",
            "HWR",
            "PAA Rule 1 — always-planned procedure CCS",
            "CCS",
        ),
        (
            "_hwr_value_set_paa2",
            "ccs_diagnosis_category",
            "HWR",
            "PAA Rule 2 — always-planned diagnosis CCS",
            "CCS",
        ),
        (
            "_hwr_value_set_paa3",
            "category_or_code",
            "HWR",
            "PAA Rule 3 — potentially-planned procedure",
            "CCS/ICD-10-PCS",
        ),
        (
            "_hwr_value_set_paa4",
            "category_or_code",
            "HWR",
            "PAA Rule 3 gate — acute diagnosis",
            "CCS/ICD-10-CM",
        ),
    ]

    # Tables where `code_type` mixes CCS categories with redundant ICD-10
    # detail expansions.  The measure logic operates at CCS level — ICD-10
    # detail rows are reference-only and should be diffed separately.
    _HAS_CODE_TYPE = {
        "_uamcc_value_set_paa3",
        "_uamcc_value_set_paa4",
        "_acr_value_set_paa3",
        "_acr_value_set_paa4",
        "_hwr_value_set_paa3",
        "_hwr_value_set_paa4",
        "_uamcc_value_set_exclusions",
    }

    def _build_code_sets(_q, _key_col, _years, _code_type_filter=None):
        """Return {year: {norm_code: orig_code}} dicts."""
        _by_year = {}
        for _y in _years:
            _where = (
                f" AND code_type = '{_code_type_filter}'" if _code_type_filter else ""
            )
            _raw = [
                r[0]
                for r in con.execute(
                    f'SELECT DISTINCT "{_key_col}" FROM {_q}'
                    f" WHERE performance_year = {_y}"
                    f' AND "{_key_col}" IS NOT NULL{_where}'
                ).fetchall()
            ]
            _by_year[_y] = {_norm(v): v for v in _raw}
        return _by_year

    def _diff_years(_by_year, _years, _measure, _role, _code_type, _tbl):
        """Diff consecutive year pairs and return code change rows."""
        _rows = []
        for _i in range(len(_years) - 1):
            _ya, _yb = _years[_i], _years[_i + 1]
            _sa = set(_by_year[_ya])
            _sb = set(_by_year[_yb])
            for _c in sorted(_sb - _sa):
                _rows.append(
                    {
                        "measure": _measure,
                        "role": _role,
                        "code_type": _code_type,
                        "code": _by_year[_yb][_c],
                        "change": "added",
                        "transition": f"{_ya} -> {_yb}",
                        "table": _tbl,
                    }
                )
            for _c in sorted(_sa - _sb):
                _rows.append(
                    {
                        "measure": _measure,
                        "role": _role,
                        "code_type": _code_type,
                        "code": _by_year[_ya][_c],
                        "change": "removed",
                        "transition": f"{_ya} -> {_yb}",
                        "table": _tbl,
                    }
                )
        return _rows

    _code_rows = []
    for _tbl, _key_col, _measure, _role, _code_type in _VS:
        _q = f'cms_quality_measures."{_tbl}"'
        try:
            _years = sorted(
                r[0]
                for r in con.execute(
                    f"SELECT DISTINCT performance_year FROM {_q} ORDER BY 1"
                ).fetchall()
            )
        except Exception:
            continue
        if len(_years) < 2:
            continue

        if _tbl in _HAS_CODE_TYPE:
            # Diff CCS-level entries (functionally meaningful)
            _ccs_sets = _build_code_sets(_q, _key_col, _years, "CCS")
            _code_rows.extend(
                _diff_years(
                    _ccs_sets,
                    _years,
                    _measure,
                    _role + " (CCS — operative)",
                    "CCS",
                    _tbl,
                )
            )
            # Diff ICD-10 detail entries separately (reference-only)
            for _icd_type in ("ICD-10-CM", "ICD-10-PCS"):
                _icd_sets = _build_code_sets(_q, _key_col, _years, _icd_type)
                if any(len(v) > 0 for v in _icd_sets.values()):
                    _code_rows.extend(
                        _diff_years(
                            _icd_sets,
                            _years,
                            _measure,
                            _role + f" ({_icd_type} — reference detail)",
                            _icd_type,
                            _tbl,
                        )
                    )
        else:
            _by_year = _build_code_sets(_q, _key_col, _years)
            _code_rows.extend(
                _diff_years(
                    _by_year,
                    _years,
                    _measure,
                    _role,
                    _code_type,
                    _tbl,
                )
            )

    _codes_df = (
        pl.DataFrame(_code_rows)
        if _code_rows
        else pl.DataFrame({"note": ["No code changes detected"]})
    )

    # Summary by measure × role × direction
    _summary = (
        _codes_df.group_by("measure", "role", "code_type", "change", "transition")
        .agg(pl.col("code").count().alias("n_codes"))
        .sort("measure", "role", "transition", "change")
    )

    # Impact classification: which changes could shift measure results?
    _impact_md = """
    | Change Type | Potential Impact | What to Query |
    |-------------|-----------------|---------------|
    | **Cohort inclusion** codes added | More patients enter the denominator | `WHERE dx_code IN ({codes}) AND encounter_type = 'acute inpatient'` |
    | **Cohort inclusion** codes removed | Fewer patients in denominator | Same query — patients with these codes drop out |
    | **Exclusion** codes added | More encounters excluded from numerator/denominator | `WHERE dx_ccs IN ({codes})` on your index admissions |
    | **Exclusion** codes removed | Fewer exclusions → larger effective denominator | Same query — previously excluded patients now included |
    | **PAA Rule 1/2/3** codes added | More admissions classified as *planned* → lower unplanned rate | `WHERE procedure_ccs IN ({codes})` or `WHERE dx_ccs IN ({codes})` on readmissions |
    | **PAA Rule 3 gate (acute dx)** codes added | More procedures remain *unplanned* (acute dx negates Rule 3) → higher unplanned rate | `WHERE dx_ccs IN ({codes})` on readmissions with potentially-planned procedures |
    | **CCS crosswalk** codes added/remapped | Diagnosis-to-CCS mapping changes cascade into all CCS-based logic above | `WHERE dx_code IN ({codes})` — check if CCS category assignment changed |
    """

    mo.vstack(
        [
            mo.md(
                """## Research Strategy — Spec Year Code Changes

    To measure the real-world impact of spec year changes, query the **specific codes
    that changed** against a target population. The tables below enumerate every code
    added or removed between consecutive spec years, tagged by measure, functional role,
    and code type.

    ### How spec changes propagate through the measures

    """
                + _impact_md
                + """

    ### Step-by-step research protocol

    1. **Export the code change table** below (CSV download via table widget)
    2. **Filter to your measure of interest** (UAMCC, ACR, or HWR)
    3. **Query your claims population** for encounters matching the changed codes:
       - For ICD-10-CM changes: join on `principal_diagnosis_code` or `condition.normalized_code`
       - For ICD-10-PCS changes: join on `procedure.normalized_code` or `hcpcs_code`
       - For CCS changes: first map your ICD codes through the CCS crosswalk, then match
    4. **Count affected encounters** — the overlap between changed codes and your population
       determines whether the spec change would shift the measure result
    5. **Re-run the pipeline** with each spec year's value sets (using the sensitivity
       analysis cell above) on your real data to quantify the actual difference
    """
            ),
            mo.md(f"### Change Summary — {len(_codes_df)} total code changes"),
            mo.ui.table(_summary, label="Changes by Measure / Role / Direction"),
            mo.md("### Full Code Change Inventory"),
            mo.ui.table(_codes_df, label="All Changed Codes (exportable)"),
        ]
    )
    return


if __name__ == "__main__":
    app.run()