stack/dev/scripts/generate_ccw_models.py

"""Generate Pydantic SQLTable models from the CCW FFS Claims record layout.

Reads ``record-layout-ffs-claims.xlsx`` and produces one Python module
per claim type (IP, SNF, Hospice, HHA, HOP, Carrier, DME), each
containing one ``SQLTable`` subclass per sub-table (Base claim,
Revenue center, Line, Condition code, etc.).

Usage:
    python generate_ccw_models.py [--xlsx PATH] [--out DIR]

Defaults:
    --xlsx  record-layout-ffs-claims.xlsx  (same directory)
    --out   ../src/ccw/table/
"""

from __future__ import annotations

import argparse
from pathlib import Path

import fastexcel

# ── Sheets and their slug names ──────────────────────────────────────────────

SHEETS = {
    "IP": "ip",
    "SNF": "snf",
    "Hospice": "hospice",
    "HHA": "hha",
    "HOP": "hop",
    "Carrier": "carrier",
    "DME": "dme",
}

# Section names that are actual sub-tables (not footnotes)
VALID_SECTIONS = {
    "Base claim file",
    "Revenue center file",
    "Condition code file",
    "Occurrence code file",
    "Span code file",
    "Value code file",
    "Demonstrations/Innovations code file",
    "Line file",
}

# ── CCW SAS type → Python type ───────────────────────────────────────────────

SAS_TYPE_MAP = {
    "CHAR": "str",
    "NUM": "float",
    "DATE": "date",
}


def python_type(sas_type: str) -> str:
    return SAS_TYPE_MAP.get(sas_type, "str")


def needs_import(py_type: str) -> set[str]:
    if py_type == "date":
        return {"from datetime import date"}
    return set()


# ── Name helpers ─────────────────────────────────────────────────────────────

SECTION_SLUG = {
    "Base claim file": "base",
    "Revenue center file": "revenue_center",
    "Condition code file": "condition_code",
    "Occurrence code file": "occurrence_code",
    "Span code file": "span_code",
    "Value code file": "value_code",
    "Demonstrations/Innovations code file": "demo",
    "Line file": "line",
}


def to_class_name(claim_type: str, section: str) -> str:
    """Build PascalCase class name from claim type + section slug."""
    slug = SECTION_SLUG.get(section, section.lower().replace(" ", "_"))
    raw = f"{claim_type}_{slug}"
    parts = raw.split("_")
    return "".join(p.capitalize() for p in parts if p)


def to_table_name(section: str) -> str:
    """Build snake_case table name from section name."""
    return SECTION_SLUG.get(section, section.lower().replace(" ", "_"))


# ── Parsing ──────────────────────────────────────────────────────────────────


def parse_sheet(
    xlsx_path: str, sheet_name: str
) -> dict[str, list[tuple[str, str, str]]]:
    """Parse a single Excel sheet into sections of fields.

    Returns a dict mapping section name to list of
    (long_sas_name, sas_type, label) tuples.
    """
    f = fastexcel.read_excel(xlsx_path)
    sheet = f.load_sheet_by_name(sheet_name, header_row=3)
    df = sheet.to_polars()

    # Standardize column names
    cols = df.columns
    df = df.rename(
        {
            cols[0]: "short_name",
            cols[1]: "long_name",
            cols[2]: "label",
            cols[3]: "type",
            cols[4]: "length",
            cols[5]: "sequence",
            cols[6]: "new_vars",
        }
    )

    sections: dict[str, list[tuple[str, str, str]]] = {}
    current_section: str | None = None

    for row in df.iter_rows(named=True):
        short = row["short_name"]
        long = row["long_name"]
        sas_type = row["type"]

        # Section header: short_name has text, long_name is null
        if long is None and short is not None:
            if short.strip() in VALID_SECTIONS:
                current_section = short.strip()
                sections[current_section] = []
            continue

        # Data row: both long_name and type are present
        if long is not None and sas_type is not None and current_section is not None:
            field_name = long.strip().lower()
            label = (row["label"] or "").strip()
            sections[current_section].append((field_name, sas_type.strip(), label))

    return sections


# ── Code generation ──────────────────────────────────────────────────────────


def generate_model(
    claim_slug: str,
    section: str,
    fields: list[tuple[str, str, str]],
) -> tuple[str, set[str]]:
    """Generate a single Pydantic model class."""
    cls = to_class_name(claim_slug, section)
    table = to_table_name(section)
    all_imports: set[str] = set()
    field_lines: list[str] = []

    for field_name, sas_type, label in fields:
        py = python_type(sas_type)
        all_imports |= needs_import(py)
        # All fields nullable — CCW data has lots of missing values
        field_lines.append(f"    {field_name}: {py} | None = None")

    lines = [
        f"class {cls}(SQLTable):",
        f'    """CCW FFS Claims: {claim_slug.upper()} / {section}"""',
        "",
        '    __schema__ = "ccw"',
        f'    __tablename__ = "{claim_slug}_{table}"',
        "",
    ]
    lines.extend(field_lines)

    return "\n".join(lines), all_imports


def generate_module(
    xlsx_path: str,
    sheet_name: str,
    claim_slug: str,
) -> str:
    """Generate a full Python module for one claim type."""
    sections = parse_sheet(xlsx_path, sheet_name)
    all_imports: set[str] = {"from __future__ import annotations"}
    all_imports.add("from aco.table.base import SQLTable")
    models: list[str] = []

    for section, fields in sections.items():
        if not fields:
            continue
        model_code, imp = generate_model(claim_slug, section, fields)
        all_imports |= imp
        models.append(model_code)

    sorted_imports = sorted(all_imports)
    header = "\n".join(sorted_imports)
    body = "\n\n\n".join(models)

    return f"{header}\n\n\n{body}\n"


# ── Main ────────────────────────────────────────────────────────────────────


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Generate Pydantic models from CCW FFS Claims record layout"
    )
    parser.add_argument(
        "--xlsx",
        default=str(
            Path(__file__).resolve().parents[2]
            / "dev"
            / "seeds"
            / "record-layout-ffs-claims.xlsx"
        ),
        help="Path to the CCW record layout Excel file",
    )
    parser.add_argument(
        "--out",
        default=str(Path(__file__).resolve().parents[2] / "src" / "ccw" / "table"),
        help="Output directory",
    )
    args = parser.parse_args()

    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    # Copy base.py reference — ccw reuses aco.table.base.SQLTable
    total_tables = 0

    for sheet_name, slug in SHEETS.items():
        code = generate_module(args.xlsx, sheet_name, slug)
        module_path = out_dir / f"{slug}.py"
        module_path.write_text(code)

        sections = parse_sheet(args.xlsx, sheet_name)
        n = sum(1 for fields in sections.values() if fields)
        total_tables += n
        print(f"Wrote {module_path} ({n} tables)")

    # Write __init__.py
    init_lines = [
        '"""CCW FFS Claims Pydantic table models."""',
        "",
    ]
    for slug in sorted(SHEETS.values()):
        init_lines.append(f"from . import {slug} as {slug}")
    init_lines.append("")
    (out_dir / "__init__.py").write_text("\n".join(init_lines))
    print(f"\nWrote {out_dir / '__init__.py'}")

    # Write ccw top-level __init__.py if it doesn't exist
    ccw_init = out_dir.parent / "__init__.py"
    if not ccw_init.exists():
        ccw_init.write_text(
            '"""CCW — CMS Chronic Conditions Warehouse data models."""\n'
        )
        print(f"Wrote {ccw_init}")

    print(f"\nDone: {len(SHEETS)} claim types, {total_tables} sub-tables → {out_dir}/")


if __name__ == "__main__":
    main()