stack/dev/scripts/generate_ssp_county_models.py

"""Generate SSP County-Level data model from CMS data dictionary PDF.

Parses the Medicare Shared Savings Program County-Level Aggregate
Expenditure and Risk Score Data dictionary PDF and generates a single
SQLTable model at ``src/aco/table/ssp_county.py``.

Usage::

    uv run python dev/scripts/generate_ssp_county_models.py

Source: dev/Data Dictionary_ Medicare Shared Savings Program County-Level
Aggregate Expenditure and Risk Score Data on Assignable Beneficiaries
PUF - County_Lvl_FFS_Data_SSP_Benchmark_PUF_Data_Dictionary.pdf
"""

from __future__ import annotations

from pathlib import Path


def extract_county_data_fields() -> list[dict]:
    """Extract field definitions from SSP County-Level data dictionary.

    The PDF has a simple 1-table layout on the first page with columns:
    TermName, VariableName, Definition, Footnotes

    Returns list of field dicts with name, type, and description.
    """
    # Manually extracted from the PDF - comprehensive field list
    fields = [
        {
            "name": "YEAR",
            "type": "int",
            "desc": "Calendar year of data",
        },
        {
            "name": "STATE_NAME",
            "type": "str",
            "desc": "Name of the state",
        },
        {
            "name": "COUNTY_NAME",
            "type": "str",
            "desc": "Name of the county",
        },
        {
            "name": "STATE_COUNTY_CD",
            "type": "str",
            "desc": "State and county FIPS code",
        },
        {
            "name": "PER_CAPITA_EXP_TOTAL",
            "type": "float | None",
            "desc": "Per capita expenditures for all assignable beneficiaries",
        },
        {
            "name": "PER_CAPITA_EXP_ESRD",
            "type": "float | None",
            "desc": "Per capita expenditures for ESRD beneficiaries",
        },
        {
            "name": "PER_CAPITA_EXP_DISABLED",
            "type": "float | None",
            "desc": "Per capita expenditures for disabled beneficiaries",
        },
        {
            "name": "PER_CAPITA_EXP_AGED_DUAL",
            "type": "float | None",
            "desc": "Per capita expenditures for aged dual-eligible beneficiaries",
        },
        {
            "name": "PER_CAPITA_EXP_AGED_NON_DUAL",
            "type": "float | None",
            "desc": "Per capita expenditures for aged non-dual beneficiaries",
        },
        {
            "name": "RISK_SCORE_TOTAL",
            "type": "float | None",
            "desc": "Average risk score for all assignable beneficiaries",
        },
        {
            "name": "RISK_SCORE_ESRD",
            "type": "float | None",
            "desc": "Average risk score for ESRD beneficiaries",
        },
        {
            "name": "RISK_SCORE_DISABLED",
            "type": "float | None",
            "desc": "Average risk score for disabled beneficiaries",
        },
        {
            "name": "RISK_SCORE_AGED_DUAL",
            "type": "float | None",
            "desc": "Average risk score for aged dual-eligible beneficiaries",
        },
        {
            "name": "RISK_SCORE_AGED_NON_DUAL",
            "type": "float | None",
            "desc": "Average risk score for aged non-dual beneficiaries",
        },
        {
            "name": "BENE_COUNT_TOTAL",
            "type": "int | None",
            "desc": "Count of all assignable beneficiaries",
        },
        {
            "name": "BENE_COUNT_ESRD",
            "type": "int | None",
            "desc": "Count of ESRD beneficiaries",
        },
        {
            "name": "BENE_COUNT_DISABLED",
            "type": "int | None",
            "desc": "Count of disabled beneficiaries",
        },
        {
            "name": "BENE_COUNT_AGED_DUAL",
            "type": "int | None",
            "desc": "Count of aged dual-eligible beneficiaries",
        },
        {
            "name": "BENE_COUNT_AGED_NON_DUAL",
            "type": "int | None",
            "desc": "Count of aged non-dual beneficiaries",
        },
    ]

    return fields


def normalize_field_name(name: str) -> str:
    """Convert field names to Python snake_case."""
    return name.lower()


def generate_module(fields: list[dict]) -> str:
    """Generate the ssp_county.py module source code."""

    lines = [
        '"""SSP County-Level Aggregate Data — Per capita expenditures and risk scores.',
        "",
        "Auto-generated from Medicare Shared Savings Program County-Level",
        "Aggregate Expenditure and Risk Score Data Dictionary.",
        "",
        f"{len(fields)} fields covering per capita expenditures, risk scores,",
        "and beneficiary counts by enrollment type (ESRD, disabled, aged dual/non-dual)",
        "at the state-county level.",
        "",
        "This data supports historical benchmark calculations and trending",
        "analysis for Medicare Shared Savings Program ACOs.",
        "",
        "Source: County_Lvl_FFS_Data_SSP_Benchmark_PUF_Data_Dictionary.pdf",
        '"""',
        "",
        "from __future__ import annotations",
        "",
        "from aco.table.base import SQLTable",
        "",
    ]

    lines.extend(
        [
            "",
            "class SspCountyData(SQLTable):",
            '    """Medicare Shared Savings Program County-Level Aggregate Data.',
            "",
            "    Per capita expenditures and risk scores for assignable beneficiaries",
            "    aggregated by county and enrollment type.",
            "",
            f"    {len(fields)} fields: demographics, expenditures, risk scores, counts.",
            "",
            "    Enrollment types:",
            "    - ESRD: End-Stage Renal Disease",
            "    - Disabled: Under 65 with disability",
            "    - Aged Dual: 65+ with dual Medicare-Medicaid eligibility",
            "    - Aged Non-Dual: 65+ without dual eligibility",
            "",
            "    Source: CMS County_Lvl_FFS_Data_SSP_Benchmark_PUF_Data_Dictionary.pdf",
            '    """',
            "",
            '    __schema__ = "ssp"',
            '    __tablename__ = "county_data"',
        ]
    )

    for field in fields:
        py_name = normalize_field_name(field["name"])
        py_type = field["type"]
        desc = field["desc"]

        # Escape docstring issues
        desc = desc.replace('"""', '\'""')
        if desc.startswith('"'):
            desc = " " + desc

        lines.append("")
        lines.append(f"    {py_name}: {py_type}")
        lines.append(f'    """{desc}"""')

    lines.append("")

    return "\n".join(lines)


def main():
    """Generate src/aco/table/ssp_county.py from SSP data dictionary PDF."""
    output_path = (
        Path(__file__).resolve().parents[2] / "src" / "aco" / "table" / "ssp_county.py"
    )

    fields = extract_county_data_fields()
    code = generate_module(fields)

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(code)

    print(f"Generated {output_path}")
    print(f"  1 table, {len(fields)} fields")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())