stack/dev/scripts/ingest_opps.py

"""Ingest CMS OPPS addendum files into DuckDB.

Parses Addendum A (APC weights), Addendum B (HCPCS→APC map), and
wage index files from downloaded ZIPs. Handles column naming
variations across 13 years of CMS format drift (CY2014–2026).

Produces:
  * opps.apc_weight    (Addendum A — APC relative weights)
  * opps.addendum_b    (Addendum B — HCPCS→APC crosswalk)
  * opps.wage_index    (Wage index by CBSA)

Usage:
    uv run python dev/scripts/download_opps_files.py   # download first
    uv run python dev/scripts/ingest_opps.py           # then ingest
    uv run python dev/scripts/ingest_opps.py --year 2026
"""

from __future__ import annotations

import argparse
import io
import re
import zipfile
from pathlib import Path

import duckdb
import pandas as pd

ROOT = Path(__file__).resolve().parents[2]
OPPS_DIR = ROOT / "data" / "cms" / "opps"
DUCKDB_PATH = ROOT / "data" / "aco.duckdb"

# ---------------------------------------------------------------------------
# Column normalisation — CMS changes headers across years
# ---------------------------------------------------------------------------

ADDENDUM_A_COL_MAP = {
    "apc": "apc",
    "apc group": "apc",
    "apc group number": "apc",
    "group title": "group_title",
    "apc group title": "group_title",
    "status indicator": "status_indicator",
    "si": "status_indicator",
    "relative weight": "relative_weight",
    "payment rate": "payment_rate",
    "national unadjusted payment": "payment_rate",
    "minimum unadjusted copayment": "minimum_unadjusted_copayment",
}

ADDENDUM_B_COL_MAP = {
    "hcpcs code": "hcpcs",
    "hcpcs": "hcpcs",
    "cpt/hcpcs": "hcpcs",
    "short descriptor": "short_description",
    "short description": "short_description",
    "si": "status_indicator",
    "status indicator": "status_indicator",
    "apc": "apc",
    "apc group": "apc",
    "apc group number": "apc",
    "apc title": "apc_title",
    "relative weight": "relative_weight",
    "payment rate": "payment_rate",
    "national unadjusted payment": "payment_rate",
    "minimum unadjusted copayment": "minimum_unadjusted_copayment",
}

WAGE_INDEX_COL_MAP = {
    "cbsa": "cbsa",
    "cbsa number": "cbsa",
    "cbsa name": "cbsa_name",
    "urban area title": "cbsa_name",
    "state": "state",
    "wage index": "wage_index",
    "pre-reclassification wage index": "wage_index",
    "reclassified wage index": "reclassified_wage_index",
    "post-reclassification wage index": "reclassified_wage_index",
}


def normalise_columns(df: pd.DataFrame, col_map: dict[str, str]) -> pd.DataFrame:
    """Lowercase columns, apply mapping, keep only mapped columns."""
    df.columns = [str(c).strip().lower() for c in df.columns]
    df = df.rename(columns=col_map)
    present = [c for c in set(col_map.values()) if c in df.columns]
    return df[present]


# ---------------------------------------------------------------------------
# Readers — handle ZIP contents (CSV or Excel inside)
# ---------------------------------------------------------------------------


def find_data_file(zf: zipfile.ZipFile, pattern: str = "") -> str | None:
    """Find a specific addendum file inside a combined addenda ZIP.

    CMS OPPS addenda ZIPs contain multiple addenda (A, B, C, D1, D2, etc.)
    as both XLSX and CSV (508 version). Prefer XLSX in the root, fall back
    to CSV in the 508 subfolder.
    """
    skip = ["__macosx", ".ds_store", "readme", "note"]
    candidates = []
    for name in zf.namelist():
        low = name.lower()
        if any(s in low for s in skip):
            continue
        if low.endswith((".csv", ".xlsx", ".xls")):
            candidates.append(name)
    if pattern:
        # Match "Addendum A" / "Addendum B" etc., separator-insensitively:
        # CMS names use spaces ("Addendum B.xlsx") and hyphens
        # ("CMS-1613-FC-Addendum-B.xlsx") interchangeably across years.
        pn = re.sub(r"[^a-z0-9]", "", pattern.lower())
        filtered = [c for c in candidates if pn in re.sub(r"[^a-z0-9]", "", c.lower())]
        # Exclude the "Data Addendum B" geometric-mean-cost file — it carries
        # SI/APC/Comment-Indicator only, no payment rate, so it never
        # reconciles. Its name varies ("Data-Addendum-B", "DataAddB",
        # "Data Add B"); normalise and drop anything containing "dataadd".
        real = [
            c for c in filtered if "dataadd" not in re.sub(r"[^a-z0-9]", "", c.lower())
        ]
        pool = real or filtered
        # Prefer xlsx, then xls, then csv; prefer root over a 508 subfolder.
        for in_root in (True, False):
            for ext in (".xlsx", ".xls", ".csv"):
                hits = [
                    f
                    for f in pool
                    if f.lower().endswith(ext) and (("/" not in f) == in_root)
                ]
                if hits:
                    return sorted(hits)[0]
        if pool:
            return sorted(pool)[0]
    return candidates[0] if candidates else None


def _find_header_row(df: pd.DataFrame, col_map: dict[str, str]) -> int:
    """Find the row index containing column headers in a preamble-laden CMS file.

    Requires >= 3 matches AND at least 3 non-null cells (to distinguish
    a real header row from a title row that happens to contain keywords).
    """
    target_cols = [k for k in col_map if k not in ("", "year")]
    for i, row in df.iterrows():
        vals = [str(v).strip().lower() for v in row if pd.notna(v) and str(v).strip()]
        if len(vals) < 3:
            continue
        matches = sum(
            1 for v in vals for k in target_cols if v == k or (len(k) > 3 and k in v)
        )
        if matches >= 3:
            return i
    return 0


def read_from_zip(
    zpath: Path, col_map: dict[str, str], pattern: str = ""
) -> pd.DataFrame:
    """Read and normalise a data file from inside a ZIP.

    Handles CMS preamble rows in both CSV and Excel formats.
    """
    try:
        with zipfile.ZipFile(zpath) as zf:
            target = find_data_file(zf, pattern)
            if not target:
                return pd.DataFrame()

            data = zf.read(target)

            if target.lower().endswith(".csv"):
                raw = data.decode("latin-1")
                lines = raw.splitlines()
                # Find header row
                hdr_idx = 0
                target_cols = list(col_map.keys())[:5]
                for i, line in enumerate(lines[:20]):
                    low = line.strip().lower()
                    if sum(1 for k in target_cols if k in low) >= 2:
                        hdr_idx = i
                        break
                buf = "\n".join(lines[hdr_idx:])
                df = pd.read_csv(io.StringIO(buf), dtype=str, on_bad_lines="skip")
            else:
                # Excel: read without header to find preamble extent
                df_raw = pd.read_excel(
                    io.BytesIO(data), dtype=str, header=None, nrows=20
                )
                hdr_row = _find_header_row(df_raw, col_map)
                df = pd.read_excel(io.BytesIO(data), dtype=str, header=hdr_row)

            return normalise_columns(df, col_map)
    except Exception as e:
        print(f"    ERROR: {zpath.name}/{pattern}: {e}")
        return pd.DataFrame()


def clean_numeric(series: pd.Series) -> pd.Series:
    """Strip $ and commas, convert to float."""
    return (
        series.astype(str)
        .str.replace("$", "", regex=False)
        .str.replace(",", "", regex=False)
        .str.strip()
        .pipe(pd.to_numeric, errors="coerce")
    )


# ---------------------------------------------------------------------------
# Main ingestion
# ---------------------------------------------------------------------------


def ingest_all(con: duckdb.DuckDBPyConnection, year_filter: str = "") -> None:
    """Ingest all downloaded OPPS files into DuckDB."""

    con.execute("CREATE SCHEMA IF NOT EXISTS opps")

    apc_frames: list[pd.DataFrame] = []
    addb_frames: list[pd.DataFrame] = []
    wage_frames: list[pd.DataFrame] = []

    year_dirs = sorted(OPPS_DIR.iterdir())
    if year_filter:
        year_dirs = [d for d in year_dirs if d.name == year_filter]

    for year_dir in year_dirs:
        if not year_dir.is_dir():
            continue
        year = year_dir.name
        print(f"\n--- CY{year} ---")

        # Combined addenda ZIP (CY2021+)
        addenda_zip = year_dir / f"opps_{year}_addenda.zip"
        if not addenda_zip.exists():
            # Try old single-file patterns
            addenda_zip = year_dir / f"opps_{year}_addendum_a.zip"

        if addenda_zip.exists():
            # Addendum A — APC weights
            df = read_from_zip(addenda_zip, ADDENDUM_A_COL_MAP, pattern="addendum a")
            if not df.empty:
                df["year"] = int(year)
                for col in [
                    "relative_weight",
                    "payment_rate",
                    "minimum_unadjusted_copayment",
                ]:
                    if col in df.columns:
                        df[col] = clean_numeric(df[col])
                apc_frames.append(df)
                print(f"    Addendum A: {len(df)} APCs")

            # Addendum B — HCPCS→APC crosswalk
            df = read_from_zip(addenda_zip, ADDENDUM_B_COL_MAP, pattern="addendum b")
            if not df.empty:
                df["year"] = int(year)
                for col in [
                    "relative_weight",
                    "payment_rate",
                    "minimum_unadjusted_copayment",
                ]:
                    if col in df.columns:
                        df[col] = clean_numeric(df[col])
                addb_frames.append(df)
                print(f"    Addendum B: {len(df)} HCPCS codes")

    # Load into DuckDB
    print("\n--- Loading into DuckDB ---")

    def _load(table: str, frames: list[pd.DataFrame]) -> None:
        """Replace `table` — but with --year, merge instead of wipe.

        A year-filtered run only parses that year's files, so a bare
        DROP/CREATE here used to erase every other year (#509). When
        filtering, keep the existing table's other-year rows and rebuild
        from their union; pandas concat also unions columns, so CMS
        format drift between years stays handled either way.
        """
        merged = pd.concat(frames, ignore_index=True)
        if year_filter:
            exists = con.execute(
                "SELECT 1 FROM information_schema.tables "
                "WHERE table_schema = 'opps' AND table_name = ?",
                [table],
            ).fetchone()
            if exists:
                other_years = con.execute(
                    f"SELECT * FROM opps.{table} WHERE year != ?",  # noqa: S608
                    [int(year_filter)],
                ).df()
                merged = pd.concat([other_years, merged], ignore_index=True)
        con.execute(f"DROP TABLE IF EXISTS opps.{table}")
        con.execute(f"CREATE TABLE opps.{table} AS SELECT * FROM merged")  # noqa: S608
        print(f"  opps.{table}: {len(merged)} rows ({merged['year'].nunique()} years)")

    if apc_frames:
        _load("apc_weight", apc_frames)

    if addb_frames:
        _load("addendum_b", addb_frames)

    if wage_frames:
        _load("wage_index", wage_frames)

    # Skin sub specific: extract skin sub codes from Addendum B history.
    # Derived from the (merged) table rather than this run's frames, so a
    # --year run keeps the other years here too.
    if addb_frames:
        print("\n--- Skin substitute history in Addendum B ---")
        con.execute("DROP TABLE IF EXISTS opps.skin_sub_addendum_b")
        con.execute(
            "CREATE TABLE opps.skin_sub_addendum_b AS "
            "SELECT * FROM opps.addendum_b "
            "WHERE regexp_matches(CAST(hcpcs AS VARCHAR), '^Q4\\d{2,3}$|^C527[1-8]$')"
        )
        n, years, codes = con.execute(
            "SELECT count(*), count(DISTINCT year), count(DISTINCT hcpcs) "
            "FROM opps.skin_sub_addendum_b"
        ).fetchone()
        print(
            f"  opps.skin_sub_addendum_b: {n} rows ({years} years, {codes} unique codes)"
        )

        # SI distribution for skin subs (column absent in some year formats)
        has_si = con.execute(
            "SELECT 1 FROM information_schema.columns "
            "WHERE table_schema = 'opps' AND table_name = 'skin_sub_addendum_b' "
            "AND column_name = 'status_indicator'"
        ).fetchone()
        si_rows = (
            con.execute(
                "SELECT year, status_indicator, count(*) "
                "FROM opps.skin_sub_addendum_b "
                "GROUP BY year, status_indicator ORDER BY year, status_indicator"
            ).fetchall()
            if has_si
            else []
        )
        if si_rows:
            print("  Status indicators for skin subs by year:")
            for yr, si, ct in si_rows:
                print(f"    CY{yr} SI={si}: {ct} codes")


def main() -> None:
    parser = argparse.ArgumentParser(description="Ingest CMS OPPS files")
    parser.add_argument("--year", help="Ingest only this year")
    parser.add_argument(
        "--no-lake",
        action="store_true",
        help="skip publishing to the DuckLake lakehouse",
    )
    args = parser.parse_args()

    print("Ingesting CMS OPPS files into DuckDB ...")
    # duckdb_batch preflights the single-writer lock (retry + name the
    # holding PID) instead of dying on a raw IOException when a notebook
    # kernel holds a connection (#508).
    from conf.connect import duckdb_batch, publish_replica

    with duckdb_batch("aco") as con:
        ingest_all(con, year_filter=args.year or "")

        # Final inventory
        print("\n--- OPPS tables ---")
        for r in con.execute("""
            SELECT table_name FROM information_schema.tables
            WHERE table_schema = 'opps' ORDER BY table_name
        """).fetchall():
            cnt = con.execute(f"SELECT count(*) FROM opps.{r[0]}").fetchone()[0]  # noqa: S608
            print(f"  opps.{r[0]:30s}: {cnt:>6} rows")

    # Refresh the notebook read replica so long-running readers see the
    # new data without ever locking this primary (#510).
    replica = publish_replica("aco")
    print(f"replica → {replica}")

    if not args.no_lake:
        print("\n--- Publishing to DuckLake ---")
        import _lake

        _lake.publish_lake(("opps",))

    print("\nDone.")


if __name__ == "__main__":
    main()