stack/dev/scripts/generate_asr_table_models.py

"""Generate ASR SQLTable models from the ALR/ASR User's Guide PDF.

Parses the Assignment Summary Report (ASR) table specifications from
the ALRASRGuide.pdf to extract all 9 ASR aggregate table layouts and
generates a single Python module at ``src/aco/table/asr.py`` with one
SQLTable subclass per table.

ASR tables contain aggregate statistics (no beneficiary-identifiable data).

Usage::

    uv run python dev/scripts/generate_asr_table_models.py

Source: dev/ALRASRGuide.pdf pages 27-39
"""

from __future__ import annotations

import re
from pathlib import Path


def extract_asr_table_specs() -> dict:
    """Extract ASR table field specifications.

    ASR tables are Excel-based aggregate summary tables with counts,
    rates, and percentages rather than individual beneficiary records.

    Returns dict mapping table_id -> {title, fields: list[dict]}
    """
    tables = {}

    # Table 2-1: Beneficiary Assignment Summary
    tables["2-1"] = {
        "title": "Beneficiary Assignment Summary",
        "fields": [
            {"name": "Category", "type": "str", "desc": "Assignment category"},
            {"name": "Count", "type": "int", "desc": "Number of beneficiaries"},
        ],
    }

    # Table 2-1A: Voluntarily Aligned Beneficiaries and Exclusions
    tables["2-1A"] = {
        "title": "Voluntarily Aligned Beneficiaries and Exclusions",
        "fields": [
            {"name": "Category", "type": "str", "desc": "Exclusion category"},
            {
                "name": "Part_I_Count",
                "type": "int | None",
                "desc": "Initial prospective assignment count",
            },
            {
                "name": "Part_II_Count",
                "type": "int | None",
                "desc": "Current assignment count (quarterly/annual)",
            },
        ],
    }

    # Table 2-1B: Claims-Based Beneficiary Assignment and Exclusions
    tables["2-1B"] = {
        "title": "Claims-Based Beneficiary Assignment and Exclusions",
        "fields": [
            {
                "name": "Category",
                "type": "str",
                "desc": "Assignment or exclusion category",
            },
            {"name": "Count", "type": "int", "desc": "Number of beneficiaries"},
            {
                "name": "Percent",
                "type": "float | None",
                "desc": "Percentage of assignable population",
            },
        ],
    }

    # Table 2-4: Demographic and Eligibility Characteristics
    tables["2-4"] = {
        "title": "Demographic and Eligibility Characteristics",
        "fields": [
            {
                "name": "Characteristic",
                "type": "str",
                "desc": "Demographic or eligibility characteristic",
            },
            {
                "name": "Category",
                "type": "str | None",
                "desc": "Category within characteristic",
            },
            {
                "name": "ESRD_Person_Years",
                "type": "float | None",
                "desc": "ESRD person years",
            },
            {
                "name": "Disabled_Person_Years",
                "type": "float | None",
                "desc": "Disabled person years",
            },
            {
                "name": "Aged_Dual_Person_Years",
                "type": "float | None",
                "desc": "Aged dual-eligible person years",
            },
            {
                "name": "Aged_Non_Dual_Person_Years",
                "type": "float | None",
                "desc": "Aged non-dual person years",
            },
            {
                "name": "Total_Person_Years",
                "type": "float | None",
                "desc": "Total person years",
            },
            {
                "name": "Percent_Total",
                "type": "float | None",
                "desc": "Percentage of total person years",
            },
            {
                "name": "Median_All_ACOs",
                "type": "float | None",
                "desc": "Median value across all ACOs",
            },
        ],
    }

    # Table 2-5: Distribution of Total Assigned Beneficiary Residence by County
    tables["2-5"] = {
        "title": "Distribution of Total Assigned Beneficiary Residence by County",
        "fields": [
            {"name": "State_Name", "type": "str", "desc": "State name"},
            {"name": "County_Name", "type": "str", "desc": "County name"},
            {"name": "ESRD_Person_Years", "type": "float", "desc": "ESRD person years"},
            {
                "name": "Disabled_Person_Years",
                "type": "float",
                "desc": "Disabled person years",
            },
            {
                "name": "Aged_Dual_Person_Years",
                "type": "float",
                "desc": "Aged dual-eligible person years",
            },
            {
                "name": "Aged_Non_Dual_Person_Years",
                "type": "float",
                "desc": "Aged non-dual person years",
            },
            {
                "name": "Total_Person_Years",
                "type": "float",
                "desc": "Total person years",
            },
            {
                "name": "Percent_Total",
                "type": "float",
                "desc": "Percentage of total assigned beneficiaries",
            },
        ],
    }

    # Table 2-5A: Distribution by County, Excluding COVID-19 Episodes
    tables["2-5A"] = {
        "title": "Distribution of Total Assigned Beneficiary Residence by County, Excluding COVID-19 Episodes",
        "fields": [
            {"name": "State_Name", "type": "str", "desc": "State name"},
            {"name": "County_Name", "type": "str", "desc": "County name"},
            {
                "name": "ESRD_Person_Years",
                "type": "float",
                "desc": "ESRD person years (excluding COVID episodes)",
            },
            {
                "name": "Disabled_Person_Years",
                "type": "float",
                "desc": "Disabled person years (excluding COVID episodes)",
            },
            {
                "name": "Aged_Dual_Person_Years",
                "type": "float",
                "desc": "Aged dual-eligible person years (excluding COVID episodes)",
            },
            {
                "name": "Aged_Non_Dual_Person_Years",
                "type": "float",
                "desc": "Aged non-dual person years (excluding COVID episodes)",
            },
            {
                "name": "Total_Person_Years",
                "type": "float",
                "desc": "Total person years (excluding COVID episodes)",
            },
            {
                "name": "Percent_Total",
                "type": "float",
                "desc": "Percentage of total assigned beneficiaries",
            },
        ],
    }

    # Table 2-6: Frequencies and Rates per 10,000 Beneficiaries by Disease Group
    tables["2-6"] = {
        "title": "Frequencies and Rates per 10,000 Beneficiaries by Disease Group (CMS-HCC)",
        "fields": [
            {"name": "HCC_Category", "type": "str", "desc": "CMS-HCC payment category"},
            {
                "name": "Beneficiary_Count",
                "type": "int",
                "desc": "Number of beneficiaries with HCC",
            },
            {
                "name": "Rate_Per_10000",
                "type": "float",
                "desc": "Rate per 10,000 beneficiaries",
            },
        ],
    }

    # Table 2-7: Prevalence of COVID-19
    tables["2-7"] = {
        "title": "Prevalence of COVID-19 Among ACO Assigned Beneficiaries",
        "fields": [
            {"name": "Category", "type": "str", "desc": "COVID-19 category"},
            {"name": "ACO_Count", "type": "int", "desc": "Count for this ACO"},
            {"name": "ACO_Percent", "type": "float", "desc": "Percentage for this ACO"},
            {
                "name": "Program_Percent",
                "type": "float",
                "desc": "Program-wide percentage",
            },
        ],
    }

    # Table 2-8: Utilization of Primary Care Services
    tables["2-8"] = {
        "title": "Utilization of Primary Care Services by Assigned Beneficiaries",
        "fields": [
            {
                "name": "Service_Category",
                "type": "str",
                "desc": "Primary care service category",
            },
            {
                "name": "Service_Count",
                "type": "int",
                "desc": "Total number of services",
            },
            {
                "name": "Rate_Per_1000_Person_Years",
                "type": "float",
                "desc": "Rate per 1,000 person years",
            },
        ],
    }

    # Table 2-9: Distribution by Indicators of Underserved Populations
    tables["2-9"] = {
        "title": "Distribution of Total Assigned Beneficiaries by Indicators of Underserved Populations",
        "fields": [
            {
                "name": "Indicator",
                "type": "str",
                "desc": "Underserved population indicator",
            },
            {
                "name": "Category",
                "type": "str | None",
                "desc": "Category within indicator",
            },
            {
                "name": "Beneficiary_Count",
                "type": "int | None",
                "desc": "Number of beneficiaries",
            },
            {
                "name": "Percent_Total",
                "type": "float | None",
                "desc": "Percentage of total assigned",
            },
            {
                "name": "Person_Years",
                "type": "float | None",
                "desc": "Person years for category",
            },
        ],
    }

    return tables


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


def class_name(table_id: str) -> str:
    """Convert table ID to class name."""
    names = {
        "2-1": "AsrBeneficiaryAssignmentSummary",
        "2-1A": "AsrVoluntaryAlignmentExclusions",
        "2-1B": "AsrClaimsBasedExclusions",
        "2-4": "AsrDemographicCharacteristics",
        "2-5": "AsrCountyDistribution",
        "2-5A": "AsrCountyDistributionExcludingCovid",
        "2-6": "AsrHccPrevalence",
        "2-7": "AsrCovid19Prevalence",
        "2-8": "AsrPrimaryCareUtilization",
        "2-9": "AsrUnderservedPopulations",
    }
    return names[table_id]


def table_name(table_id: str) -> str:
    """Convert table ID to database table name."""
    cls = class_name(table_id)
    return re.sub(r"(?<!^)(?=[A-Z])", "_", cls).lower()


def generate_module(tables: dict) -> str:
    """Generate the asr.py module source code."""
    total_fields = sum(len(t["fields"]) for t in tables.values())

    lines = [
        '"""ASR — Assignment Summary Report table layouts.',
        "",
        "Auto-generated from the ALR/ASR User's Guide:",
        "dev/ALRASRGuide.pdf",
        f"Version #16 (April 2024) — {total_fields} fields across {len(tables)} tables.",
        "",
        "Assignment Summary Reports are Excel files (.xlsx) delivered to all",
        "ACOs in the Medicare Shared Savings Program. They contain aggregate",
        "statistics on assigned beneficiary populations (no beneficiary-identifiable data).",
        "",
        "Tables::",
        "",
    ]

    for tid in sorted(tables.keys()):
        info = tables[tid]
        lines.append(f"    Table {tid}: {info['title']} ({len(info['fields'])} fields)")

    lines.extend(
        [
            '"""',
            "",
            "from __future__ import annotations",
            "",
            "from aco.table.base import SQLTable",
            "",
        ]
    )

    # Generate classes in order
    for tid in sorted(tables.keys()):
        info = tables[tid]
        cls = class_name(tid)
        tbl = table_name(tid)
        title = info["title"]

        lines.append("")
        lines.append(f"class {cls}(SQLTable):")
        lines.append(f'    """ASR Table {tid}: {title}')
        lines.append("")
        lines.append(f"    {len(info['fields'])} fields, Excel (.xlsx) format.")
        lines.append("    Aggregate statistics, no beneficiary-identifiable data.")
        lines.append("")
        lines.append("    Source: ALRASRGuide.pdf")
        lines.append('    """')
        lines.append("")
        lines.append('    __schema__ = "asr"')
        lines.append(f'    __tablename__ = "{tbl}"')

        for field in info["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("")

    # Remove trailing blank lines
    while lines and lines[-1] == "":
        lines.pop()

    return "\n".join(lines) + "\n"


def main():
    """Generate src/aco/table/asr.py from ALRASRGuide.pdf."""
    output_path = (
        Path(__file__).resolve().parents[2] / "src" / "aco" / "table" / "asr.py"
    )

    tables = extract_asr_table_specs()
    code = generate_module(tables)

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

    print(f"Generated {output_path}")
    print(
        f"  {len(tables)} tables, {sum(len(t['fields']) for t in tables.values())} fields"
    )

    return 0


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