stack/dev/scripts/populate_quality_measure_value_sets.py

"""Populate CMS quality measure value set tables for all available performance years.

Reads Excel value set workbooks from Zotero storage, parses each sheet into
the corresponding SQLTable model columns, stamps each row with ``performance_year``,
and upserts the data into DuckDB.

Also produces year-over-year diff reports showing which codes were added,
removed, or unchanged between adjacent years for each value set — giving a
quantitative answer to "how different is the PY2026 spec from PY2025?"

Usage::

    # Load all available years, write to notebooks/aco.duckdb
    uv run python dev/scripts/populate_quality_measure_value_sets.py

    # Dry-run: discover files and show what would be loaded, no DB writes
    uv run python dev/scripts/populate_quality_measure_value_sets.py --dry-run

    # Load only a specific performance year
    uv run python dev/scripts/populate_quality_measure_value_sets.py --year 2025

    # Show year-over-year diffs after loading
    uv run python dev/scripts/populate_quality_measure_value_sets.py --diff

    # Write diff summary to a Parquet file in addition to console output
    uv run python dev/scripts/populate_quality_measure_value_sets.py --diff --diff-out diffs.parquet

    # Replace existing data instead of appending (full reload)
    uv run python dev/scripts/populate_quality_measure_value_sets.py --mode replace

    # Use a different Zotero storage root
    uv run python dev/scripts/populate_quality_measure_value_sets.py --zotero /path/to/zotero

Design
------
Each Excel workbook contains multiple named sheets.  Each sheet maps to one
``_value_set_*`` table in the ``cms_quality_measures`` schema.  The loader:

  1. Normalises Excel column headers to snake_case so they match SQLTable
     field names (e.g. "ICD-10-CM" → "icd_10_cm").
  2. Drops any Excel columns that have no corresponding model field.
  3. Adds the ``performance_year`` discriminator column.
  4. Writes the resulting polars DataFrame to DuckDB via INSERT … SELECT.

Year-over-year diffs compare the set of *code values* in the key column(s)
of each value set between adjacent performance years.  The summary table
reports: codes_in_a, codes_in_b, added, removed, unchanged, pct_changed.
"""

from __future__ import annotations

import argparse
import re
import textwrap
import zipfile
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import polars as pl

# ── Path constants ────────────────────────────────────────────────────────────
from conf import ROOT as _ROOT
from conf import path as _conf_path

_SEEDS = _ROOT / "dev" / "seeds"
_ZOTERO_DEFAULT = _conf_path("storage.zotero")
_DB_DEFAULT = _conf_path("db.aco")

# ── Value set spec ─────────────────────────────────────────────────────────────


@dataclass(frozen=True)
class SheetSpec:
    """Describes a single Excel sheet → DuckDB table mapping.

    Parameters
    ----------
    sheet_name : str
        Exact sheet name (or substring prefix) in the workbook.
    table_name : str
        Qualified ``schema.table`` target (leading underscore preserved).
    key_cols : tuple[str, ...]
        Column name(s) in the normalised frame that constitute the
        primary "code" for diff comparisons.  Typically one column.
    col_renames : dict[str, str]
        Optional explicit column renames applied *after* auto-normalisation.
        Use for headers that don't normalise cleanly to the model field name.
    """

    sheet_name: str
    table_name: str
    key_cols: tuple[str, ...]
    col_renames: dict[str, str] = field(default_factory=dict)


@dataclass(frozen=True)
class MifSpec:
    """One performance year × measure combination with its source files.

    Parameters
    ----------
    year : int
        Performance year (e.g. 2025, 2026).
    measure : str
        Short measure code: "UAMCC", "ACR", "HWR", "MCC".
    program : str
        Payment program: "REACH" or "MIPS".
    xlsx_basename : str | None
        Excel value set workbook filename (or None if unavailable).
    zip_basename : str | None
        Zip archive containing the xlsx (or None if direct).
    sheets : tuple[SheetSpec, ...]
        Ordered sheet specs for this workbook.
    """

    year: int
    measure: str
    program: str
    xlsx_basename: str | None
    zip_basename: str | None
    sheets: tuple[SheetSpec, ...]


# ── Sheet specs per measure ───────────────────────────────────────────────────

_UAMCC_SHEETS: tuple[SheetSpec, ...] = (
    SheetSpec(
        sheet_name="UAMCC Cohort",
        table_name="cms_quality_measures._uamcc_value_set_cohort",
        key_cols=("icd_10_cm",),
        col_renames={
            "chronic_condition": "chronic_condition_group",
            "code": "icd_10_cm",
            "number_type_of_claims_to_quality": "claims_to_qualify",
            "lookback_year": "lookback_years",
        },
    ),
    SheetSpec(
        sheet_name="UAMCC Exclusions",
        table_name="cms_quality_measures._uamcc_value_set_exclusions",
        key_cols=("category_or_code",),
        col_renames={
            "ccs_category_or_icd_10_code": "code_type",
        },
    ),
    SheetSpec(
        sheet_name="UAMCC PAA1",
        table_name="cms_quality_measures._uamcc_value_set_paa1",
        key_cols=("ccs_procedure_category",),
        col_renames={},
    ),
    SheetSpec(
        sheet_name="UAMCC PAA2",
        table_name="cms_quality_measures._uamcc_value_set_paa2",
        key_cols=("ccs_diagnosis_category",),
        col_renames={
            "ccs_description": "description",
        },
    ),
    SheetSpec(
        sheet_name="UAMCC PAA3",
        table_name="cms_quality_measures._uamcc_value_set_paa3",
        key_cols=("category_or_code",),
        col_renames={
            "ccs_procedure_category_or_icd_10_pcs_code": "code_type",
        },
    ),
    SheetSpec(
        sheet_name="UAMCC PAA4",
        table_name="cms_quality_measures._uamcc_value_set_paa4",
        key_cols=("category_or_code",),
        col_renames={
            "ccs_diagnosis_category_or_icd_10_cm_code": "code_type",
        },
    ),
    SheetSpec(
        sheet_name="UAMCC CCS-ICD10CM",
        table_name="cms_quality_measures._uamcc_value_set_ccs_icd10_cm",
        key_cols=("icd_10_cm",),
        col_renames={},
    ),
    SheetSpec(
        sheet_name="UAMCC CCS-ICD10PCS",
        table_name="cms_quality_measures._uamcc_value_set_ccs_icd10_pcs",
        key_cols=("icd_10_pcs",),
        col_renames={},
    ),
)

_ACR_SHEETS: tuple[SheetSpec, ...] = (
    SheetSpec(
        sheet_name="ACR Cohort CCS",
        table_name="cms_quality_measures._acr_value_set_cohort_ccs",
        key_cols=("ccs_category",),
        col_renames={},
    ),
    SheetSpec(
        sheet_name="ACR Cohort ICD-10",
        table_name="cms_quality_measures._acr_value_set_cohort_icd10",
        key_cols=("icd_10_pcs",),
        col_renames={},
    ),
    SheetSpec(
        sheet_name="ACR Exclusions",
        table_name="cms_quality_measures._acr_value_set_exclusions",
        key_cols=("ccs_diagnosis_category",),
        col_renames={},
    ),
    SheetSpec(
        sheet_name="ACR PAA PA1",
        table_name="cms_quality_measures._acr_value_set_paa1",
        key_cols=("ccs_procedure_category",),
        col_renames={},
    ),
    SheetSpec(
        sheet_name="ACR PAA PA2",
        table_name="cms_quality_measures._acr_value_set_paa2",
        key_cols=("ccs_diagnosis_category",),
        col_renames={
            "ccs_description": "description",
        },
    ),
    SheetSpec(
        sheet_name="ACR PAA PA3",
        table_name="cms_quality_measures._acr_value_set_paa3",
        key_cols=("category_or_code",),
        col_renames={
            "ccs_procedure_category_or_icd_10_pcs_code": "code_type",
        },
    ),
    SheetSpec(
        sheet_name="ACR PAA PA4",
        table_name="cms_quality_measures._acr_value_set_paa4",
        key_cols=("category_or_code",),
        col_renames={
            "ccs_diagnosis_category_or_icd_10_cm_code": "code_type",
        },
    ),
    # ACR CCS crosswalk sheets (present in PY2025 workbook)
    SheetSpec(
        sheet_name="ACR CCS-ICD-10CM",
        table_name="cms_quality_measures._acr_value_set_ccs_icd10_cm",
        key_cols=("icd_10_cm",),
        col_renames={},
    ),
    SheetSpec(
        sheet_name="ACR CCS-ICD-10PCS",
        table_name="cms_quality_measures._acr_value_set_ccs_icd10_pcs",
        key_cols=("icd_10_pcs",),
        col_renames={},
    ),
)

_HWR_SHEETS: tuple[SheetSpec, ...] = (
    SheetSpec(
        sheet_name="1.",  # prefix: "1. HWR Specialty Cohort Incls"
        table_name="cms_quality_measures._hwr_value_set_specialty_cohort",
        key_cols=("ccs_category",),
        col_renames={},
    ),
    SheetSpec(
        sheet_name="2.",  # prefix: "2. HWR Surg_Gyn Cohort Incls"
        table_name="cms_quality_measures._hwr_value_set_surg_gyn_cohort",
        key_cols=("icd_10_pcs",),
        col_renames={},
    ),
    SheetSpec(
        sheet_name="3.",  # prefix: "3. HWR Cohort Exclusions"
        table_name="cms_quality_measures._hwr_value_set_cohort_exclusions",
        key_cols=("ccs_diagnosis_category",),
        col_renames={},
    ),
    SheetSpec(
        sheet_name="PR.1",  # prefix: "PR.1 Always Planned Px"
        table_name="cms_quality_measures._hwr_value_set_paa1",
        key_cols=("ccs_procedure_category",),
        col_renames={},
    ),
    SheetSpec(
        sheet_name="PR.2",  # prefix: "PR.2 Always Planned Dx"
        table_name="cms_quality_measures._hwr_value_set_paa2",
        key_cols=("ccs_diagnosis_category",),
        col_renames={
            "ccs_description": "description",
        },
    ),
    SheetSpec(
        sheet_name="PR.3",  # prefix: "PR.3 Pot Planned Px"
        table_name="cms_quality_measures._hwr_value_set_paa3",
        key_cols=("category_or_code",),
        col_renames={
            "ccs_procedure_category_or_icd_10_pcs_code": "code_type",
        },
    ),
    SheetSpec(
        sheet_name="PR.4",  # prefix: "PR.4 Acute Dx"
        table_name="cms_quality_measures._hwr_value_set_paa4",
        key_cols=("category_or_code",),
        col_renames={
            "ccs_diagnosis_category_or_icd_10_cm_code": "code_type",
        },
    ),
)

# MCC MIPS uses the same measure family as UAMCC — maps to UAMCC tables.
# Sheet names in MIPS MCC workbooks use "MIPSMCC" prefix with numeric prefixes.
_MCC_SHEETS: tuple[SheetSpec, ...] = (
    SheetSpec(
        sheet_name="1. MIPSMCC",  # prefix: "1. MIPSMCC Cohort"
        table_name="cms_quality_measures._uamcc_value_set_cohort",
        key_cols=("icd_10_cm",),
        col_renames={
            "chronic_condition": "chronic_condition_group",
            "code": "icd_10_cm",
            "number_type_of_claims_to_quality": "claims_to_qualify",
            "lookback_year": "lookback_years",
        },
    ),
    SheetSpec(
        sheet_name="2. MIPSMCC",  # prefix: "2. MIPSMCC Outcome Exclusions"
        table_name="cms_quality_measures._uamcc_value_set_exclusions",
        key_cols=("category_or_code",),
        col_renames={
            "ccs_category_or_icd_10_code": "code_type",
        },
    ),
    SheetSpec(
        sheet_name="MIPSMCC PAA1",  # prefix: "MIPSMCC PAA1 Always Planned Px"
        table_name="cms_quality_measures._uamcc_value_set_paa1",
        key_cols=("ccs_procedure_category",),
        col_renames={},
    ),
    SheetSpec(
        sheet_name="MIPSMCC PAA2",  # prefix: "MIPSMCC PAA2 Always Planned Dx"
        table_name="cms_quality_measures._uamcc_value_set_paa2",
        key_cols=("ccs_diagnosis_category",),
        col_renames={
            "ccs_description": "description",
        },
    ),
    SheetSpec(
        sheet_name="MIPSMCC PAA3",  # prefix: "MIPSMCC PAA3 Pot Planned_Px"
        table_name="cms_quality_measures._uamcc_value_set_paa3",
        key_cols=("category_or_code",),
        col_renames={
            "ccs_procedure_category_or_icd_10_pcs_code": "code_type",
        },
    ),
    SheetSpec(
        sheet_name="MIPSMCC PAA4",  # prefix: "MIPSMCC PAA4 Acute Diagnoses"
        table_name="cms_quality_measures._uamcc_value_set_paa4",
        key_cols=("category_or_code",),
        col_renames={
            "ccs_diagnosis_category_or_icd_10_cm_code": "code_type",
        },
    ),
)

# ── MIF registry ──────────────────────────────────────────────────────────────

REGISTRY: tuple[MifSpec, ...] = (
    # ── MIPS HWR 2024 ────────────────────────────────────────────────
    MifSpec(
        year=2024,
        measure="HWR",
        program="MIPS",
        xlsx_basename="MIPS_PY2024_Hospital_Wide_Readmission_CodeTables_12072023.xlsx",
        zip_basename="MIPS_Hospital-Wide Readmission_2024_MIF.zip",
        sheets=_HWR_SHEETS,
    ),
    # ── ACO REACH PY2024 ─────────────────────────────────────────────
    MifSpec(
        year=2024,
        measure="UAMCC",
        program="REACH",
        xlsx_basename="ACOREACH_PY2024_UAMCC_ValueSet_updated05082024.xlsx",
        zip_basename="ACOREACHPY2024MIFsandValueSets_updated10282024.zip",
        sheets=_UAMCC_SHEETS,
    ),
    MifSpec(
        year=2024,
        measure="ACR",
        program="REACH",
        xlsx_basename="ACOREACH_PY2024_ACR_ValueSet_updated0162024.xlsx",
        zip_basename="ACOREACHPY2024MIFsandValueSets_updated10282024.zip",
        sheets=_ACR_SHEETS,
    ),
    # ── MIPS MCC 2025 ────────────────────────────────────────────────
    # MCC loads BEFORE REACH UAMCC because both map to _uamcc_value_set_*
    # tables.  REACH is the canonical source for REACH ACOs, so it loads
    # last and its upsert replaces any MCC data for the same year.
    MifSpec(
        year=2025,
        measure="MCC",
        program="MIPS",
        # File dated 01.03.25 (Jan 3, 2025) contains PY2025 value sets
        xlsx_basename="Version_2024_MIPS_MCC_01.03.25.xlsx",
        zip_basename="2025-MIPS-MCC-Measure-Specification.zip",
        sheets=_MCC_SHEETS,
    ),
    # ── MIPS MCC 2026 ────────────────────────────────────────────────
    MifSpec(
        year=2026,
        measure="MCC",
        program="MIPS",
        xlsx_basename="Version_2025_MIPS_MCC_01.09.26.xlsx",
        zip_basename="2026-MIPS-MCC-Measure-Specification.zip",
        sheets=_MCC_SHEETS,
    ),
    # ── MIPS HWR 2025 ────────────────────────────────────────────────
    MifSpec(
        year=2025,
        measure="HWR",
        program="MIPS",
        xlsx_basename="MIPS_CY_2025_HWR_Measure_Code_Specifications.xlsx",
        zip_basename="MIPS_Hospital-Wide-Readmission_2025.zip",
        sheets=_HWR_SHEETS,
    ),
    # ── ACO REACH PY2025 ─────────────────────────────────────────────
    # REACH loads after MIPS MCC — REACH upsert replaces MCC data,
    # ensuring the REACH-specific value sets are canonical.
    MifSpec(
        year=2025,
        measure="UAMCC",
        program="REACH",
        xlsx_basename="ACOREACH_PY2025_UAMCC_ValueSet_posted07072025.xlsx",
        zip_basename="ACOREACHPY2025MIFs&ValueSets_UpdatedPosting_07072025.zip",
        sheets=_UAMCC_SHEETS,
    ),
    MifSpec(
        year=2025,
        measure="ACR",
        program="REACH",
        xlsx_basename="ACOREACH_PY2025_ACR_ValueSet_posted02212025.xlsx",
        zip_basename="ACOREACHPY2025MIFs&ValueSets_UpdatedPosting_07072025.zip",
        sheets=_ACR_SHEETS,
    ),
    # ── ACO REACH PY2026 ─────────────────────────────────────────────
    MifSpec(
        year=2026,
        measure="UAMCC",
        program="REACH",
        xlsx_basename="ACOREACH-PY2026-UAMCC-ValueSet_posted10312025.xlsx",
        zip_basename="PY 2026 MIFs and Value Sets_Updated 4i Posting_12-15-2025.zip",
        sheets=_UAMCC_SHEETS,
    ),
    MifSpec(
        year=2026,
        measure="ACR",
        program="REACH",
        xlsx_basename="ACOREACH-PY2026-ACR-ValueSet_posted10312025.xlsx",
        zip_basename="PY 2026 MIFs and Value Sets_Updated 4i Posting_12-15-2025.zip",
        sheets=_ACR_SHEETS,
    ),
)

# ── Column normalisation ──────────────────────────────────────────────────────

# Known Excel → model field aliases used globally across workbooks.
# Applied *before* SheetSpec.col_renames (which are spec-specific overrides).
_GLOBAL_ALIASES: dict[str, str] = {
    # ICD code columns
    "icd10cm": "icd_10_cm",
    "icd10pcs": "icd_10_pcs",
    "icd_10_cm_code": "icd_10_cm",
    "icd_10_pcs_code": "icd_10_pcs",
    "icd10_cm": "icd_10_cm",
    "icd10_pcs": "icd_10_pcs",
    # CCS columns
    "ccs_cat": "ccs_category",
    "ccs_cat_number": "ccs_category",
    "ccs_description_label": "ccs_description",
    # Description variants
    "code_description": "description",
    "label_description": "description",
    "icd_10_cm_description": "description",
    "icd_10_pcs_description": "description",
    # Cohort columns
    "condition_group": "chronic_condition_group",
    # Category / code columns
    "category_code": "category_or_code",
    # Procedure/diagnosis type columns
    "procedure_or_diagnosis_category": "procedure_or_diagnosis",
    "outcome_exclusion_category": "exclusion_category",
    # Associated CCS columns (parenthetical stripped)
    "associated_ccs_procedure_category": "associated_ccs_category",
    "associated_ccs_diagnosis_category": "associated_ccs_category",
}


def _normalise_header(raw: str) -> str:
    """Normalise an Excel column header to a Python snake_case field name.

    Steps:
      1. Strip whitespace
      2. Remove parenthetical qualifiers — e.g. "(index claim)"
      3. Remove trailing comma-phrases — e.g. ", as of October 1, 2022"
      4. Lower and strip trailing question marks / whitespace
      5. Replace non-alphanumeric runs with underscores
      6. Collapse consecutive underscores
      7. Apply global alias table
    """
    s = raw.strip()
    s = re.sub(r"\s*\([^)]*\)", "", s)
    s = re.sub(r",\s+as of\b.*$", "", s, flags=re.IGNORECASE)
    s = s.lower().rstrip("? ")
    s = re.sub(r"[^a-z0-9]+", "_", s)
    s = re.sub(r"_+", "_", s).strip("_")
    return _GLOBAL_ALIASES.get(s, s)


# ── File discovery helpers ────────────────────────────────────────────────────


def _find_in_zotero(basename: str, zotero_root: Path) -> Path | None:
    """Locate a file anywhere under the Zotero storage tree by basename."""
    for candidate in zotero_root.rglob(basename):
        if candidate.is_file():
            return candidate
    return None


def _extract_from_zip(
    zip_basename: str,
    member_basename: str,
    zotero_root: Path,
    dest_dir: Path,
) -> Path | None:
    """Extract a single member from a Zotero zip archive and return its path."""
    zip_path = _find_in_zotero(zip_basename, zotero_root)
    if zip_path is None:
        return None
    dest = dest_dir / member_basename
    if dest.exists():
        return dest
    try:
        with zipfile.ZipFile(zip_path) as zf:
            for info in zf.infolist():
                if Path(info.filename).name == member_basename:
                    with zf.open(info) as src, dest.open("wb") as out:
                        out.write(src.read())
                    return dest
    except Exception as exc:  # noqa: BLE001
        print(f"  ⚠  Could not extract {member_basename} from {zip_path.name}: {exc}")
    return None


def _resolve_xlsx(
    spec: MifSpec,
    zotero_root: Path,
    work_dir: Path,
) -> Path | None:
    """Return path to the Excel workbook for a MifSpec, extracting from zip if needed."""
    if spec.xlsx_basename is None:
        return None
    direct = _find_in_zotero(spec.xlsx_basename, zotero_root)
    if direct:
        return direct
    if spec.zip_basename:
        return _extract_from_zip(
            spec.zip_basename, spec.xlsx_basename, zotero_root, work_dir
        )
    return None


# ── Excel parsing ─────────────────────────────────────────────────────────────


def _match_sheet(available: list[str], spec_name: str) -> str | None:
    """Find the best matching sheet name for a SheetSpec.sheet_name.

    Matching order:
      1. Exact match
      2. Case-insensitive exact
      3. Available sheet starts with spec_name (prefix match)
    """
    if spec_name in available:
        return spec_name
    lower_spec = spec_name.lower()
    for name in available:
        if name.lower() == lower_spec:
            return name
    for name in available:
        if name.lower().startswith(lower_spec.lower()):
            return name
    return None


def _parse_sheet(
    wb: Any,  # openpyxl Workbook
    sheet_name: str,
    spec: SheetSpec,
    performance_year: int,
) -> "pl.DataFrame | None":
    """Parse one worksheet into a polars DataFrame with normalised columns.

    Returns None if the sheet is empty or all-null.
    """
    import polars as pl

    ws = wb[sheet_name]
    rows = list(ws.iter_rows(values_only=True))

    # Find header row: first row with ≥2 non-None cells
    header_row_idx: int | None = None
    for i, row in enumerate(rows[:5]):
        non_none = [c for c in row if c is not None]
        if len(non_none) >= 2:
            header_row_idx = i
            break

    if header_row_idx is None:
        return None

    raw_headers = [
        str(c).strip() if c is not None else f"_col{j}"
        for j, c in enumerate(rows[header_row_idx])
    ]
    normalised_headers = [_normalise_header(h) for h in raw_headers]

    # Apply spec-specific renames
    for old, new in spec.col_renames.items():
        normalised_headers = [new if h == old else h for h in normalised_headers]

    data_rows = rows[header_row_idx + 1 :]
    if not data_rows:
        return None

    # Build column dict — skip all-None columns
    col_dict: dict[str, list[Any]] = {h: [] for h in normalised_headers}
    for row in data_rows:
        # Pad or truncate row to match header length
        padded = list(row) + [None] * max(0, len(normalised_headers) - len(row))
        for h, val in zip(normalised_headers, padded[: len(normalised_headers)]):
            col_dict[h].append(val)

    # Drop completely empty columns and duplicate-name columns (keep first)
    seen: set[str] = set()
    clean: dict[str, list[Any]] = {}
    for col, values in col_dict.items():
        if col in seen:
            continue
        seen.add(col)
        if any(v is not None for v in values):
            clean[col] = values

    if not clean:
        return None

    # Drop rows that are entirely None (trailing blank rows in Excel)
    n_rows = len(next(iter(clean.values())))
    keep = [any(clean[c][i] is not None for c in clean) for i in range(n_rows)]
    clean = {c: [v for v, k in zip(vals, keep) if k] for c, vals in clean.items()}

    if not any(clean.values()):
        return None

    # Cast to polars, coercing all values to string (value sets
    # are reference tables — string representation avoids type mismatches)
    str_clean: dict[str, list[str | None]] = {
        c: [str(v).strip() if v is not None else None for v in vals]
        for c, vals in clean.items()
    }
    df = pl.DataFrame(str_clean)

    # Add performance_year column
    df = df.with_columns(pl.lit(performance_year).alias("performance_year"))

    # ── Data cleanup ──────────────────────────────────────────────────────
    # CMS Excel workbooks contain artifacts that are not actual value set data.

    # 1. Drop "end of worksheet" sentinel rows (present in some PY2026 sheets).
    text_cols = [c for c in df.columns if df[c].dtype == pl.Utf8]
    if text_cols:
        sentinel_mask = pl.lit(False)
        for c in text_cols:
            sentinel_mask = sentinel_mask | pl.col(c).str.to_lowercase().str.contains(
                "end of worksheet"
            )
        df = df.filter(~sentinel_mask)

    # 2. Drop section header rows where all key columns are null.
    #    Excel workbooks use ALL CAPS group headers (e.g. "ACUTE MYOCARDIAL
    #    INFARCTION") as visual separators — they have no code values.
    key_cols_present = [kc for kc in spec.key_cols if kc in df.columns]
    if key_cols_present:
        all_keys_null = pl.lit(True)
        for kc in key_cols_present:
            all_keys_null = all_keys_null & pl.col(kc).is_null()
        df = df.filter(~all_keys_null)

    # 3. Normalise ICD-10 codes: strip dots for consistency across years.
    #    PY2024 workbooks use dotted format (I21.01), PY2025+ use undotted
    #    (I2101).  Undotted is the canonical CMS billing format.
    for col in df.columns:
        if col in ("icd_10_cm", "icd_10_pcs"):
            df = df.with_columns(pl.col(col).str.replace_all(r"\.", ""))

    return df


# ── Model field introspection ────────────────────────────────────────────────


def _model_columns(table_name: str) -> list[str] | None:
    """Return the field names for the SQLTable model matching *table_name*.

    *table_name* is ``schema.tablename`` (e.g.
    ``cms_quality_measures._uamcc_value_set_cohort``).
    Returns None if no matching model is found.
    """
    from aco.table import cms_quality_measures as tbl_mod
    from aco.table.base import SQLTable

    for attr in dir(tbl_mod):
        cls = getattr(tbl_mod, attr)
        if (
            isinstance(cls, type)
            and issubclass(cls, SQLTable)
            and cls is not SQLTable
            and f"{cls.__schema__}.{cls.__tablename__}" == table_name
        ):
            return cls.column_names()
    return None


# ── DuckDB writing ──────────────────────────────────────────────────────────


def _ensure_table(con: Any, table_name: str, columns: list[str]) -> None:
    """Create schema and table if they don't already exist."""
    schema, tbl = table_name.split(".", 1)
    qualified = f'"{schema}"."{tbl}"'
    con.execute(f'CREATE SCHEMA IF NOT EXISTS "{schema}"')
    try:
        con.execute(f"SELECT 1 FROM {qualified} LIMIT 0")
    except Exception:
        col_defs = []
        for c in columns:
            if c == "performance_year":
                col_defs.append(f'"{c}" INTEGER')
            else:
                col_defs.append(f'"{c}" VARCHAR')
        con.execute(f"CREATE TABLE {qualified} ({', '.join(col_defs)})")


def _delete_year(con: Any, table_name: str, year: int) -> int:
    """Delete rows for a specific performance year.  Returns count deleted."""
    schema, tbl = table_name.split(".", 1)
    qualified = f'"{schema}"."{tbl}"'
    try:
        result = con.execute(
            f"SELECT count(*) FROM {qualified} WHERE performance_year = {year}"
        ).fetchone()
        n = result[0] if result else 0
        if n > 0:
            con.execute(f"DELETE FROM {qualified} WHERE performance_year = {year}")
        return n
    except Exception:
        return 0


def _insert_df(con: Any, table_name: str, df: "pl.DataFrame") -> int:
    """Insert a polars DataFrame into the target table.  Returns row count."""
    schema, tbl = table_name.split(".", 1)
    qualified = f'"{schema}"."{tbl}"'

    # Cast performance_year to int for DuckDB
    if "performance_year" in df.columns:
        df = df.with_columns(pl.col("performance_year").cast(pl.Int64))

    col_list = ", ".join(f'"{c}"' for c in df.columns)
    con.execute(f"INSERT INTO {qualified} ({col_list}) SELECT {col_list} FROM df")
    return len(df)


def _write_sheet(
    con: Any,
    df: "pl.DataFrame",
    spec: SheetSpec,
    performance_year: int,
    *,
    mode: str = "upsert",
) -> int:
    """Write a parsed sheet DataFrame to DuckDB.

    Parameters
    ----------
    mode : str
        ``"upsert"`` — delete existing rows for this year then insert.
        ``"replace"`` — drop and recreate the table, then insert.
        ``"append"`` — insert without deleting.

    Returns row count written.
    """
    model_cols = _model_columns(spec.table_name)
    if model_cols is None:
        print(f"    ⚠  No model found for {spec.table_name} — skipping")
        return 0

    # Filter DataFrame to model columns only
    keep = [c for c in df.columns if c in model_cols]
    df = df.select(keep)

    # Ensure performance_year is present
    if "performance_year" not in df.columns:
        df = df.with_columns(pl.lit(performance_year).alias("performance_year"))

    _ensure_table(con, spec.table_name, model_cols)

    schema, tbl = spec.table_name.split(".", 1)
    qualified = f'"{schema}"."{tbl}"'

    if mode == "replace":
        con.execute(f"DROP TABLE IF EXISTS {qualified}")
        _ensure_table(con, spec.table_name, model_cols)
    elif mode == "upsert":
        n = _delete_year(con, spec.table_name, performance_year)
        if n:
            print(f"      (deleted {n:,} existing PY{performance_year} rows)")

    return _insert_df(con, spec.table_name, df)


# ── Workbook loader ──────────────────────────────────────────────────────────


def _load_workbook(
    mif: MifSpec,
    zotero_root: Path,
    work_dir: Path,
    con: Any | None,
    *,
    mode: str = "upsert",
    dry_run: bool = False,
) -> dict[str, int]:
    """Parse and load all sheets for one MifSpec.

    Returns ``{table_name: rows_written}`` dict.
    """
    import openpyxl

    xlsx_path = _resolve_xlsx(mif, zotero_root, work_dir)
    if xlsx_path is None:
        print(
            f"  ⚠  {mif.program} {mif.measure} PY{mif.year}: "
            f"workbook not found ({mif.xlsx_basename})"
        )
        return {}

    print(f"  📖 {mif.program} {mif.measure} PY{mif.year}: {xlsx_path.name}")

    wb = openpyxl.load_workbook(str(xlsx_path), read_only=True, data_only=True)
    available_sheets = wb.sheetnames
    results: dict[str, int] = {}

    for sheet_spec in mif.sheets:
        matched = _match_sheet(available_sheets, sheet_spec.sheet_name)
        if matched is None:
            print(
                f"    ⚠  Sheet '{sheet_spec.sheet_name}' not found "
                f"(available: {available_sheets})"
            )
            continue

        df = _parse_sheet(wb, matched, sheet_spec, mif.year)
        if df is None or len(df) == 0:
            print(f"    ⚠  Sheet '{matched}' is empty — skipping")
            continue

        if dry_run:
            print(
                f"    ✓  {matched} → {sheet_spec.table_name}: "
                f"{len(df)} rows, cols={df.columns}"
            )
            results[sheet_spec.table_name] = len(df)
        else:
            rows = _write_sheet(con, df, sheet_spec, mif.year, mode=mode)
            print(f"    ✓  {matched} → {sheet_spec.table_name}: {rows:,} rows")
            results[sheet_spec.table_name] = rows

    wb.close()
    return results


# ── Year-over-year diff ──────────────────────────────────────────────────────


@dataclass
class DiffRow:
    """One row in the year-over-year diff summary."""

    table_name: str
    key_cols: tuple[str, ...]
    year_a: int
    year_b: int
    codes_in_a: int
    codes_in_b: int
    added: int
    removed: int
    unchanged: int

    @property
    def pct_changed(self) -> float:
        total = self.codes_in_a + self.added
        if total == 0:
            return 0.0
        return round((self.added + self.removed) / total * 100, 1)


def _diff_table(
    con: Any,
    table_name: str,
    key_cols: tuple[str, ...],
    year_a: int,
    year_b: int,
) -> DiffRow | None:
    """Compare key column values between two performance years."""
    schema, tbl = table_name.split(".", 1)
    qualified = f'"{schema}"."{tbl}"'

    key_expr = " || '|' || ".join(f"COALESCE(\"{k}\", '')" for k in key_cols)

    try:
        rows_a = con.execute(
            f"SELECT DISTINCT {key_expr} AS k FROM {qualified} "
            f"WHERE performance_year = {year_a}"
        ).fetchall()
        rows_b = con.execute(
            f"SELECT DISTINCT {key_expr} AS k FROM {qualified} "
            f"WHERE performance_year = {year_b}"
        ).fetchall()
    except Exception:
        return None

    set_a = {r[0] for r in rows_a}
    set_b = {r[0] for r in rows_b}

    return DiffRow(
        table_name=table_name,
        key_cols=key_cols,
        year_a=year_a,
        year_b=year_b,
        codes_in_a=len(set_a),
        codes_in_b=len(set_b),
        added=len(set_b - set_a),
        removed=len(set_a - set_b),
        unchanged=len(set_a & set_b),
    )


def _compute_diffs(con: Any) -> list[DiffRow]:
    """Compute year-over-year diffs for all value set tables with ≥2 years."""
    table_keys: dict[str, tuple[str, ...]] = {}
    for mif in REGISTRY:
        for sheet in mif.sheets:
            if sheet.table_name not in table_keys:
                table_keys[sheet.table_name] = sheet.key_cols

    diffs: list[DiffRow] = []
    for table_name, key_cols in sorted(table_keys.items()):
        schema, tbl = table_name.split(".", 1)
        qualified = f'"{schema}"."{tbl}"'
        try:
            years = [
                r[0]
                for r in con.execute(
                    f"SELECT DISTINCT performance_year FROM {qualified} "
                    f"ORDER BY performance_year"
                ).fetchall()
            ]
        except Exception:
            continue
        for i in range(len(years) - 1):
            row = _diff_table(con, table_name, key_cols, years[i], years[i + 1])
            if row is not None:
                diffs.append(row)

    return diffs


def _print_diffs(diffs: list[DiffRow]) -> None:
    """Pretty-print the diff summary table."""
    if not diffs:
        print("\n  No year-over-year diffs available (need ≥2 years loaded).\n")
        return

    print("\n  Year-over-year value set diffs:")
    print("  " + "─" * 100)
    print(
        f"  {'Table':<50} {'Years':<12} {'In A':>6} {'In B':>6} "
        f"{'Added':>6} {'Removed':>7} {'Same':>6} {'%Chg':>6}"
    )
    print("  " + "─" * 100)
    for d in diffs:
        short = d.table_name.split(".", 1)[1] if "." in d.table_name else d.table_name
        print(
            f"  {short:<50} {d.year_a}→{d.year_b}  "
            f"{d.codes_in_a:>6} {d.codes_in_b:>6} "
            f"{d.added:>6} {d.removed:>7} "
            f"{d.unchanged:>6} {d.pct_changed:>5.1f}%"
        )
    print("  " + "─" * 100)


def _write_diff_parquet(diffs: list[DiffRow], path: Path) -> None:
    """Write diff summary to Parquet."""
    rows = [
        {
            "table_name": d.table_name,
            "key_cols": "|".join(d.key_cols),
            "year_a": d.year_a,
            "year_b": d.year_b,
            "codes_in_a": d.codes_in_a,
            "codes_in_b": d.codes_in_b,
            "added": d.added,
            "removed": d.removed,
            "unchanged": d.unchanged,
            "pct_changed": d.pct_changed,
        }
        for d in diffs
    ]
    pl.DataFrame(rows).write_parquet(str(path))
    print(f"\n  Diff summary written to {path}")


# ── CLI ──────────────────────────────────────────────────────────────────────


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Load CMS quality measure value sets into DuckDB.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=textwrap.dedent("""\
            examples:
              uv run python dev/scripts/populate_quality_measure_value_sets.py
              uv run python dev/scripts/populate_quality_measure_value_sets.py --dry-run
              uv run python dev/scripts/populate_quality_measure_value_sets.py --year 2025
              uv run python dev/scripts/populate_quality_measure_value_sets.py --diff
              uv run python dev/scripts/populate_quality_measure_value_sets.py --mode replace
        """),
    )
    parser.add_argument(
        "--zotero",
        type=Path,
        default=_ZOTERO_DEFAULT,
        help="Zotero storage root (default: %(default)s)",
    )
    parser.add_argument(
        "--db",
        type=Path,
        default=_DB_DEFAULT,
        help="DuckDB database path (default: %(default)s)",
    )
    parser.add_argument(
        "--year",
        type=int,
        default=None,
        help="Load only a specific performance year",
    )
    parser.add_argument(
        "--mode",
        choices=["upsert", "replace", "append"],
        default="upsert",
        help="Write mode: upsert (default), replace, or append",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Discover files and show what would be loaded, no DB writes",
    )
    parser.add_argument(
        "--diff",
        action="store_true",
        help="Show year-over-year diffs after loading",
    )
    parser.add_argument(
        "--diff-out",
        type=Path,
        default=None,
        help="Write diff summary to a Parquet file",
    )
    args = parser.parse_args()

    # Validate Zotero path
    if not args.zotero.is_dir():
        parser.error(f"Zotero storage not found: {args.zotero}")

    # Work directory for extracted files
    work_dir = _SEEDS / ".value_set_work"
    work_dir.mkdir(exist_ok=True)

    # Filter registry
    specs = REGISTRY
    if args.year is not None:
        specs = tuple(m for m in specs if m.year == args.year)
        if not specs:
            parser.error(f"No registry entries for year {args.year}")

    print(f"\n  Loading {len(specs)} value set workbook(s)…\n")

    # Open DuckDB connection (unless dry-run)
    con = None
    if not args.dry_run:
        import duckdb

        con = duckdb.connect(str(args.db))

    total_rows = 0
    total_tables = 0

    for mif in specs:
        results = _load_workbook(
            mif,
            args.zotero,
            work_dir,
            con,
            mode=args.mode,
            dry_run=args.dry_run,
        )
        for _tbl, rows in results.items():
            total_tables += 1
            total_rows += rows

    print(f"\n  Done: {total_rows:,} total rows across {total_tables} table loads.\n")

    # Year-over-year diffs
    if args.diff and con is not None:
        diffs = _compute_diffs(con)
        _print_diffs(diffs)
        if args.diff_out:
            _write_diff_parquet(diffs, args.diff_out)

    if con is not None:
        con.close()


if __name__ == "__main__":
    main()