stack/dev/scripts/ingest_asp.py

"""Ingest CMS Average Sales Price (ASP) Drug Pricing Files.

Reads all quarterly ASP ZIP files from data/cms/asp/, normalises column
names across 20+ years of CMS format drift, filters **strictly** to
skin substitute HCPCS codes defined in data/cms/skin_substitute_hcpcs.csv,
and produces:

  * data/cms/asp/skin_subs_asp_quarterly.csv   (flat file with provenance)
  * skin_subs.asp_quarterly                     (DuckDB table)

Every output row carries ``source_url`` and ``source_file`` columns for
full provenance — each observation traces back to a specific CMS
quarterly ZIP and the file within it that was parsed.

Source URLs are read from data/cms/asp_source_urls.csv (canonical URLs
provided from CMS.gov, one per quarter 2005-Q1 through 2026-Q2).

Column mapping rationale
------------------------
CMS has changed column names across 20+ years of ASP files:

  2005–2006: "HCPCS Code", "Short Description", "HCPCS Code Dosage",
             "Estimated ASP", "Payment Limit"
  2007–2011: Same but sometimes "Dosage" instead of "HCPCS Code Dosage"
  2012–2019: "HCPCS Code", "Short Description", "Dosage",
             "Payment Limit"
  2020–2025: "HCPCS Code", "Short Description", "HCPCS Code Dosage",
             "Payment Limit"
  2026+:     "HCPCS Code", "Short Description", "Dosage or Unit",
             "Payment Limit" (now "Payment Limit" for skin subs is
             flat $127.14/cm² under reclassification)

All variants are mapped to: hcpcs_code, short_description,
dosage_or_unit, payment_limit.

The ``asp_per_unit`` column is derived: payment_limit / 1.06, since
CMS payment limit = ASP + 6% for drugs/biologicals. For 2026+ skin
subs under reclassification, this derivation no longer applies (flat
rate), but we retain it for consistency and flag it in notes.

Filtering
---------
We filter STRICTLY to codes present in skin_substitute_hcpcs.csv
(the code universe built from CMS rulemaking in issue #232). This
is critical because the Q4xxx prefix includes non-skin-sub codes
(Q4054 darbepoetin alfa, Q4055 epoetin alfa, Q4074-Q4082 various
ESRD drugs) that must be excluded.

Usage:
    uv run python dev/scripts/ingest_asp.py
"""

from __future__ import annotations

import csv
import io
import re
import zipfile
from pathlib import Path

import pandas as pd

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
ROOT = Path(__file__).resolve().parents[2]
ASP_DIR = ROOT / "data" / "cms" / "asp"
HCPCS_REF = ROOT / "data" / "cms" / "skin_substitute_hcpcs.csv"
URL_REF = ROOT / "data" / "cms" / "asp_source_urls.csv"
DUCKDB_PATH = ROOT / "data" / "aco.duckdb"
OUTPUT_CSV = ASP_DIR / "skin_subs_asp_quarterly.csv"


def _load_source_urls() -> dict[str, str]:
    """Load canonical source URLs: {quarter: url}."""
    urls = {}
    if URL_REF.exists():
        with URL_REF.open() as f:
            for row in csv.DictReader(f):
                urls[row["quarter"]] = row["url"]
    return urls


# ---------------------------------------------------------------------------
# Column normalisation
# ---------------------------------------------------------------------------
# Mapping from observed CMS column names (lowercased) to our standard names.
# Each entry documents which CMS file years use that variant.
COL_MAP = {
    # HCPCS code — consistent across all years
    "hcpcs code": "hcpcs_code",
    "hcpcs_code": "hcpcs_code",
    # Short description — consistent
    "short description": "short_description",
    "short_description": "short_description",
    # Dosage/unit — varies by year
    "hcpcs code dosage": "dosage_or_unit",  # 2005-2006, 2020-2025
    "hcpcs_code_dosage": "dosage_or_unit",
    "dosage": "dosage_or_unit",  # 2007-2019
    "dosage or unit": "dosage_or_unit",  # 2026+
    # Payment limit — consistent (ASP + 6%)
    "payment limit": "payment_limit",
    "payment_limit": "payment_limit",
}

KEEP_COLS = ["hcpcs_code", "short_description", "dosage_or_unit", "payment_limit"]


def normalise_columns(df: pd.DataFrame) -> pd.DataFrame:
    """Lower-case columns, apply COL_MAP, keep only KEEP_COLS."""
    df.columns = [str(c).strip().lower() for c in df.columns]
    df = df.rename(columns=COL_MAP)
    present = [c for c in KEEP_COLS if c in df.columns]
    return df[present]


# ---------------------------------------------------------------------------
# Readers
# ---------------------------------------------------------------------------


def _find_header_row(lines: list[str]) -> int:
    """Return the 0-based row index of the column-header row.

    The header row starts with 'HCPCS Code,' (possibly quoted).
    Must NOT match note lines that incidentally mention 'HCPCS code'.
    """
    for i, line in enumerate(lines):
        stripped = line.strip().lower()
        if stripped.startswith("hcpcs code,") or stripped.startswith('"hcpcs code",'):
            return i
    return -1


def read_csv_from_zip(zf: zipfile.ZipFile, name: str) -> pd.DataFrame:
    """Read a section-508 CSV inside a ZIP, skipping preamble rows."""
    raw = zf.read(name).decode("latin-1")
    lines = raw.splitlines()
    hdr = _find_header_row(lines)
    if hdr < 0:
        return pd.DataFrame()
    buf = "\n".join(lines[hdr:])
    df = pd.read_csv(io.StringIO(buf), dtype=str, on_bad_lines="skip")
    return normalise_columns(df)


def read_xls_from_zip(zf: zipfile.ZipFile, name: str) -> pd.DataFrame:
    """Read .xls inside a ZIP, skipping preamble rows."""
    import xlrd

    data = zf.read(name)
    wb = xlrd.open_workbook(file_contents=data)
    sh = wb.sheet_by_index(0)
    hdr_row = -1
    for r in range(min(30, sh.nrows)):
        first_cell = str(sh.cell_value(r, 0)).strip().lower()
        if first_cell == "hcpcs code":
            hdr_row = r
            break
    if hdr_row < 0:
        return pd.DataFrame()
    headers = [str(sh.cell_value(hdr_row, c)).strip() for c in range(sh.ncols)]
    rows = []
    for r in range(hdr_row + 1, sh.nrows):
        row = [str(sh.cell_value(r, c)).strip() for c in range(sh.ncols)]
        if any(row):
            rows.append(row)
    df = pd.DataFrame(rows, columns=headers)
    return normalise_columns(df)


def pick_pricing_file(names: list[str]) -> str | None:
    """From a list of files in a ZIP, pick the ASP pricing file.

    Prefer CSV (section 508) over XLS. Skip NOC, crosswalk, NDC,
    and "not payable" files.
    """
    skip_patterns = ["noc", "crosswalk", "ndc", "not payable", "not_payable"]
    candidates: list[str] = []
    for n in names:
        low = n.lower()
        if any(s in low for s in skip_patterns):
            continue
        if low.endswith((".csv", ".xls", ".xlsx")):
            candidates.append(n)

    csvs = [c for c in candidates if c.lower().endswith(".csv")]
    if csvs:
        for c in csvs:
            low = c.lower()
            if "payment limit" in low or "pricing" in low:
                return c
        return csvs[0]

    xlss = [c for c in candidates if c.lower().endswith((".xls", ".xlsx"))]
    if xlss:
        for c in xlss:
            low = c.lower()
            if "payment limit" in low or "pricing" in low or "asp" in low:
                return c
        return xlss[0]

    return None


# ---------------------------------------------------------------------------
# Numeric cleaning
# ---------------------------------------------------------------------------


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


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main() -> None:
    # Load skin-substitute HCPCS codes (strict filter — no regex fallback)
    hcpcs_ref = pd.read_csv(HCPCS_REF, dtype=str)
    skin_codes = set(hcpcs_ref["hcpcs_code"].str.strip().str.upper())
    print(f"Loaded {len(skin_codes)} skin-substitute HCPCS codes from reference")

    # Load canonical source URLs
    source_urls = _load_source_urls()
    print(f"Loaded {len(source_urls)} source URLs")

    # Ingest all quarters
    zips = sorted(ASP_DIR.glob("asp_*.zip"))
    print(f"Found {len(zips)} quarterly ZIP files\n")

    frames: list[pd.DataFrame] = []
    for zpath in zips:
        m = re.match(r"asp_(\d{4}-Q\d)", zpath.name)
        if not m:
            print(f"  SKIP {zpath.name} (cannot parse quarter)")
            continue
        quarter = m.group(1)

        try:
            with zipfile.ZipFile(zpath) as zf:
                target = pick_pricing_file(zf.namelist())
                if target is None:
                    print(f"  SKIP {quarter}: no pricing file found in ZIP")
                    continue

                if target.lower().endswith(".csv"):
                    df = read_csv_from_zip(zf, target)
                else:
                    df = read_xls_from_zip(zf, target)

                if df.empty:
                    print(f"  SKIP {quarter}: empty after parsing {target}")
                    continue

                df["quarter"] = quarter
                df["source_file"] = target
                df["source_url"] = source_urls.get(quarter, "")
                frames.append(df)
                print(f"  OK   {quarter}: {len(df):>5} rows from {target}")
        except Exception as e:
            print(f"  ERR  {quarter}: {e}")

    if not frames:
        raise RuntimeError("No data parsed from any file")

    all_df = pd.concat(frames, ignore_index=True)
    print(f"\nTotal rows across all quarters: {len(all_df)}")

    # Normalise HCPCS codes
    all_df["hcpcs_code"] = all_df["hcpcs_code"].astype(str).str.strip().str.upper()

    # Filter STRICTLY to skin substitute code universe
    # NO regex fallback — only codes in skin_substitute_hcpcs.csv
    skin_df = all_df[all_df["hcpcs_code"].isin(skin_codes)].copy()
    print(f"Skin substitute rows (strict filter): {len(skin_df)}")

    # Show what we excluded
    q4_but_not_skin = all_df[
        all_df["hcpcs_code"].str.match(r"^Q4\d{2,3}$")
        & ~all_df["hcpcs_code"].isin(skin_codes)
    ]["hcpcs_code"].unique()
    if len(q4_but_not_skin) > 0:
        print(
            f"  Excluded Q4 codes NOT in skin sub universe: {sorted(q4_but_not_skin)}"
        )

    if not skin_df.empty:
        skin_df["payment_limit"] = clean_numeric(skin_df["payment_limit"])
        # ASP per unit = payment_limit / 1.06
        # (CMS payment limit = ASP + 6% for drugs/biologicals)
        skin_df["asp_per_unit"] = (skin_df["payment_limit"] / 1.06).round(3)

    # Reorder columns
    out_cols = [
        "quarter",
        "hcpcs_code",
        "short_description",
        "asp_per_unit",
        "payment_limit",
        "dosage_or_unit",
        "source_url",
        "source_file",
    ]
    for c in out_cols:
        if c not in skin_df.columns:
            skin_df[c] = None
    skin_df = skin_df[out_cols].sort_values(["quarter", "hcpcs_code"])

    # Write CSV
    skin_df.to_csv(OUTPUT_CSV, index=False)
    print(f"\nWrote {len(skin_df)} rows to {OUTPUT_CSV}")

    if len(skin_df) > 0:
        quarters = sorted(skin_df["quarter"].unique())
        codes = sorted(skin_df["hcpcs_code"].unique())
        print(f"  Quarters: {len(quarters)} ({quarters[0]} to {quarters[-1]})")
        print(f"  Unique HCPCS codes: {len(codes)}")

    # Load into DuckDB
    print(f"\nLoading into DuckDB at {DUCKDB_PATH} ...")
    # Lock preflight + retry instead of a raw IOException when a notebook
    # kernel holds the single-writer file (#508).
    from conf.connect import duckdb_batch, publish_replica

    with duckdb_batch("aco") as con:
        con.execute("CREATE SCHEMA IF NOT EXISTS skin_subs")
        con.execute("DROP TABLE IF EXISTS skin_subs.asp_quarterly")
        if not skin_df.empty:
            con.execute(
                """
                CREATE TABLE skin_subs.asp_quarterly AS
                SELECT * FROM read_csv_auto(?, header=true)
                """,
                [str(OUTPUT_CSV)],
            )
            count = con.execute(
                "SELECT count(*) FROM skin_subs.asp_quarterly"
            ).fetchone()[0]
            print(f"  Loaded {count} rows into skin_subs.asp_quarterly")
        else:
            print("  WARNING: No skin substitute data found in ASP files")
            print(
                "  This may be expected — skin subs may not appear in main ASP pricing files"
            )
            print("  They may be in separate NOC or tissue coding files")

    # Refresh the notebook read replica (#510).
    print(f"replica → {publish_replica('aco')}")


if __name__ == "__main__":
    main()