#!/usr/bin/env python3
"""Stream a CDC natality fixed-width file from ZIP and aggregate Robson metrics.

The script intentionally retains only aggregate counts. It never extracts or writes
individual birth records.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import subprocess
import sys
import time
import zipfile
from collections import Counter, defaultdict
from pathlib import Path


# Fixed-width positions are zero-based here. The corresponding public-use guide
# positions are one-based and stable across the 2016-2024 files.
POS = {
    "year": (8, 12),
    "facility": 49,
    "age9": 78,
    "residence": 103,
    "race_ethnicity": 116,
    "education": 123,
    "live_birth_order": 178,
    "smoking": 268,
    "bmi_recode": 286,
    "pregest_diabetes": 312,
    "gest_diabetes": 313,
    "pregest_hypertension": 314,
    "gest_hypertension": 315,
    "eclampsia": 316,
    "previous_preterm": 317,
    "previous_cesareans": (331, 333),
    "induction": 382,
    "presentation": 400,
    "trial_labor": 402,
    "delivery_method": 407,
    "apgar5": (443, 445),
    "plurality": 453,
    "gestation": (498, 500),
    "birthweight": (503, 507),
    "ventilation_6h": 517,
    "nicu": 518,
}


AGE_LABELS = {
    "1": "<15",
    "2": "15-19",
    "3": "20-24",
    "4": "25-29",
    "5": "30-34",
    "6": "35-39",
    "7": "40-44",
    "8": "45-49",
    "9": "50-54",
}

RACE_LABELS = {
    "1": "Non-Hispanic White",
    "2": "Non-Hispanic Black",
    "3": "Non-Hispanic AIAN",
    "4": "Non-Hispanic Asian",
    "5": "Non-Hispanic NHOPI",
    "6": "Non-Hispanic multiracial",
    "7": "Hispanic",
    "8": "Unknown",
}

BMI_LABELS = {
    "1": "Underweight",
    "2": "Normal weight",
    "3": "Overweight",
    "4": "Obesity class I",
    "5": "Obesity class II",
    "6": "Obesity class III",
    "9": "Unknown",
}


def byte_char(line: bytes, pos: int) -> str:
    if pos >= len(line):
        return ""
    return chr(line[pos])


def byte_int(line: bytes, bounds: tuple[int, int]) -> int | None:
    raw = line[bounds[0] : bounds[1]]
    if not raw or not raw.isdigit():
        return None
    return int(raw)


def parity_proxy(live_birth_order: int | None, previous_cesareans: int | None) -> str | None:
    """Return nulliparous/multiparous using live-birth order plus uterine scar.

    Prior cesarean overrides a first live-birth order because it documents a prior
    delivery even when the earlier infant was not live-born.
    """

    if previous_cesareans is not None and previous_cesareans > 0:
        return "multiparous"
    if live_birth_order == 1:
        return "nulliparous"
    if live_birth_order is not None and 2 <= live_birth_order <= 8:
        return "multiparous"
    return None


def labor_onset(induction: str, delivery: str, trial_labor: str) -> tuple[str | None, str | None]:
    """Classify spontaneous, induced, or prelabor cesarean onset.

    Returns (onset, data-quality flag). Delivery is "1" vaginal or "2" cesarean.
    """

    if induction == "Y":
        conflict = "induction_yes_trial_no" if delivery == "2" and trial_labor == "N" else None
        return "induced", conflict
    if induction != "N":
        return None, "induction_unknown"
    if delivery == "1":
        return "spontaneous", None
    if delivery == "2" and trial_labor == "Y":
        return "spontaneous", None
    if delivery == "2" and trial_labor == "N":
        return "prelabor_cesarean", None
    return None, "trial_labor_unknown"


def robson_group(
    plurality: int | None,
    presentation: str,
    gestation: int | None,
    live_birth_order: int | None,
    previous_cesareans: int | None,
    induction: str,
    delivery: str,
    trial_labor: str,
) -> tuple[str | None, str | None, str | None]:
    """Return (Robson group, subgroup, unclassifiable reason)."""

    if plurality is None or plurality not in {1, 2, 3, 4}:
        return None, None, "plurality_unknown"
    if plurality > 1:
        return "8", None, None

    if presentation == "3":
        return "9", None, None

    parity = parity_proxy(live_birth_order, previous_cesareans)
    if presentation == "2":
        if parity is None:
            return None, None, "parity_unknown"
        return ("6", None, None) if parity == "nulliparous" else ("7", None, None)
    if presentation != "1":
        return None, None, "presentation_unknown"

    if gestation is None or not 17 <= gestation <= 47:
        return None, None, "gestation_unknown"
    if gestation < 37:
        return "10", None, None

    if parity is None:
        return None, None, "parity_unknown"
    if previous_cesareans is None:
        return None, None, "previous_cesarean_unknown"
    if previous_cesareans > 0:
        return "5", None, None

    onset, onset_issue = labor_onset(induction, delivery, trial_labor)
    if onset is None:
        return None, None, onset_issue

    if parity == "nulliparous":
        if onset == "spontaneous":
            return "1", None, None
        if onset == "induced":
            return "2", "2a", None
        return "2", "2b", None

    if onset == "spontaneous":
        return "3", None, None
    if onset == "induced":
        return "4", "4a", None
    return "4", "4b", None


def age_broad(age9: str) -> str:
    if age9 in {"1", "2", "3"}:
        return "<25"
    if age9 in {"4", "5"}:
        return "25-34"
    if age9 in {"6", "7", "8", "9"}:
        return ">=35"
    return "Unknown"


def bmi_broad(bmi: str) -> str:
    if bmi == "1":
        return "Underweight"
    if bmi == "2":
        return "Normal weight"
    if bmi == "3":
        return "Overweight"
    if bmi in {"4", "5", "6"}:
        return "Obesity"
    return "Unknown"


def yn_state(*values: str) -> str:
    if any(value == "Y" for value in values):
        return "Yes"
    if values and all(value == "N" for value in values):
        return "No"
    return "Unknown"


def bump_metric(store: defaultdict[str, Counter], key: str, cesarean: bool) -> None:
    store[key]["n"] += 1
    if cesarean:
        store[key]["cesarean"] += 1


def serialize_nested(store: defaultdict[str, Counter]) -> dict[str, dict[str, int]]:
    return {key: dict(value) for key, value in sorted(store.items())}


def file_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def aggregate(zip_path: Path, requested_year: int) -> dict:
    flow = Counter()
    missing = Counter()
    conflicts = Counter()
    characteristics = Counter()
    groups: defaultdict[str, Counter] = defaultdict(Counter)
    group_subgroups: defaultdict[str, Counter] = defaultdict(Counter)
    modifiers: defaultdict[str, Counter] = defaultdict(Counter)
    group_modifiers: defaultdict[str, Counter] = defaultdict(Counter)
    record_lengths = Counter()

    started = time.time()
    with zipfile.ZipFile(zip_path) as archive:
        members = [item for item in archive.infolist() if not item.is_dir()]
        if len(members) != 1:
            raise RuntimeError(f"Expected one data member, found {len(members)}")
        member = members[0]

    # CDC currently uses Deflate64 for some annual archives. Python's stdlib
    # zipfile can inspect that central directory but cannot decompress Deflate64,
    # whereas the system Info-ZIP implementation can stream it safely.
    process = subprocess.Popen(
        ["unzip", "-p", str(zip_path), member.filename],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    assert process.stdout is not None
    try:
        handle = process.stdout
        try:
            for line in handle:
                flow["raw_records"] += 1
                if flow["raw_records"] <= 1000:
                    record_lengths[len(line.rstrip(b"\r\n"))] += 1

                if flow["raw_records"] % 1_000_000 == 0:
                    elapsed = time.time() - started
                    print(
                        f"{requested_year}: {flow['raw_records']:,} records in {elapsed:.1f}s",
                        file=sys.stderr,
                        flush=True,
                    )

                residence = byte_char(line, POS["residence"])
                if residence not in {"1", "2", "3"}:
                    flow["excluded_non_us_resident"] += 1
                    continue
                flow["us_resident"] += 1

                facility = byte_char(line, POS["facility"])
                if facility != "1":
                    flow["excluded_nonhospital_or_unknown"] += 1
                    continue
                flow["hospital_us_resident"] += 1

                delivery = byte_char(line, POS["delivery_method"])
                if delivery not in {"1", "2"}:
                    flow["excluded_delivery_unknown"] += 1
                    continue
                flow["known_delivery"] += 1
                cesarean = delivery == "2"

                live_birth_order = byte_int(line, (178, 179))
                if live_birth_order == 9:
                    live_birth_order = None
                previous_cesareans = byte_int(line, POS["previous_cesareans"])
                if previous_cesareans == 99:
                    previous_cesareans = None
                plurality = byte_int(line, (453, 454))
                presentation = byte_char(line, POS["presentation"])
                gestation = byte_int(line, POS["gestation"])
                if gestation == 99:
                    gestation = None
                induction = byte_char(line, POS["induction"])
                trial_labor = byte_char(line, POS["trial_labor"])

                if live_birth_order is None:
                    missing["live_birth_order"] += 1
                if previous_cesareans is None:
                    missing["previous_cesareans"] += 1
                if presentation not in {"1", "2", "3"}:
                    missing["presentation"] += 1
                if gestation is None:
                    missing["gestation"] += 1
                if induction not in {"Y", "N"}:
                    missing["induction"] += 1
                if cesarean and trial_labor not in {"Y", "N"}:
                    missing["trial_labor_among_cesareans"] += 1
                if (
                    live_birth_order == 1
                    and previous_cesareans is not None
                    and previous_cesareans > 0
                ):
                    conflicts["first_live_birth_with_previous_cesarean"] += 1
                if cesarean and induction == "Y" and trial_labor == "N":
                    conflicts["induction_yes_trial_no"] += 1

                group, subgroup, reason = robson_group(
                    plurality=plurality,
                    presentation=presentation,
                    gestation=gestation,
                    live_birth_order=live_birth_order,
                    previous_cesareans=previous_cesareans,
                    induction=induction,
                    delivery=delivery,
                    trial_labor=trial_labor,
                )
                if group is None:
                    flow["excluded_unclassifiable_robson"] += 1
                    missing[f"unclassifiable_{reason or 'other'}"] += 1
                    continue

                flow["robson_classified"] += 1
                groups[group]["n"] += 1
                if cesarean:
                    groups[group]["cesarean"] += 1
                if subgroup:
                    group_subgroups[subgroup]["n"] += 1
                    if cesarean:
                        group_subgroups[subgroup]["cesarean"] += 1

                apgar5 = byte_int(line, POS["apgar5"])
                if apgar5 is not None and 0 <= apgar5 <= 10:
                    groups[group]["apgar5_known"] += 1
                    if apgar5 <= 3:
                        groups[group]["apgar5_0_3"] += 1
                nicu = byte_char(line, POS["nicu"])
                if nicu in {"Y", "N"}:
                    groups[group]["nicu_known"] += 1
                    if nicu == "Y":
                        groups[group]["nicu"] += 1
                vent6 = byte_char(line, POS["ventilation_6h"])
                if vent6 in {"Y", "N"}:
                    groups[group]["vent6_known"] += 1
                    if vent6 == "Y":
                        groups[group]["vent6"] += 1

                age9 = byte_char(line, POS["age9"])
                bmi = byte_char(line, POS["bmi_recode"])
                race = byte_char(line, POS["race_ethnicity"])
                gdm = yn_state(byte_char(line, POS["gest_diabetes"]))
                hdp = yn_state(
                    byte_char(line, POS["pregest_hypertension"]),
                    byte_char(line, POS["gest_hypertension"]),
                    byte_char(line, POS["eclampsia"]),
                )
                prior_cs = (
                    "Yes"
                    if previous_cesareans is not None and previous_cesareans > 0
                    else "No"
                    if previous_cesareans == 0
                    else "Unknown"
                )
                preterm = "Yes" if gestation is not None and gestation < 37 else "No"
                plurality_label = "Multiple" if plurality and plurality > 1 else "Singleton"
                presentation_label = {
                    "1": "Cephalic",
                    "2": "Breech",
                    "3": "Other",
                }.get(presentation, "Unknown")
                smoking = yn_state(byte_char(line, POS["smoking"]))

                characteristic_items = {
                    "age": AGE_LABELS.get(age9, "Unknown"),
                    "age_broad": age_broad(age9),
                    "bmi": BMI_LABELS.get(bmi, "Unknown"),
                    "bmi_broad": bmi_broad(bmi),
                    "race_ethnicity": RACE_LABELS.get(race, "Unknown"),
                    "gdm": gdm,
                    "hdp": hdp,
                    "previous_cesarean": prior_cs,
                    "preterm": preterm,
                    "plurality": plurality_label,
                    "presentation": presentation_label,
                    "smoking": smoking,
                }
                for name, category in characteristic_items.items():
                    characteristics[f"{name}|{category}"] += 1

                modifier_items = {
                    "age": age_broad(age9),
                    "bmi": bmi_broad(bmi),
                    "gdm": gdm,
                    "hdp": hdp,
                }
                for name, category in modifier_items.items():
                    bump_metric(modifiers, f"{name}|{category}", cesarean)
                    bump_metric(group_modifiers, f"{name}|{category}|{group}", cesarean)
        finally:
            handle.close()
        stderr = process.stderr.read().decode("utf-8", errors="replace") if process.stderr else ""
        return_code = process.wait()
        if return_code != 0:
            raise RuntimeError(f"unzip failed with exit code {return_code}: {stderr.strip()}")
    finally:
        if process.poll() is None:
            process.terminate()
            process.wait()

    return {
        "year": requested_year,
        "source_zip": zip_path.name,
        "source_zip_sha256": file_sha256(zip_path),
        "source_member": member.filename,
        "source_member_uncompressed_bytes": member.file_size,
        "source_member_crc": member.CRC,
        "record_length_sample": dict(record_lengths),
        "elapsed_seconds": round(time.time() - started, 3),
        "flow": dict(flow),
        "missing": dict(missing),
        "conflicts": dict(conflicts),
        "characteristics": dict(sorted(characteristics.items())),
        "groups": serialize_nested(groups),
        "group_subgroups": serialize_nested(group_subgroups),
        "modifiers": serialize_nested(modifiers),
        "group_modifiers": serialize_nested(group_modifiers),
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--year", type=int, required=True)
    parser.add_argument("--zip", dest="zip_path", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    result = aggregate(args.zip_path, args.year)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8") as handle:
        json.dump(result, handle, ensure_ascii=False, indent=2, sort_keys=True)
        handle.write("\n")

    print(
        json.dumps(
            {
                "year": result["year"],
                "raw_records": result["flow"].get("raw_records", 0),
                "classified": result["flow"].get("robson_classified", 0),
                "elapsed_seconds": result["elapsed_seconds"],
            }
        )
    )


if __name__ == "__main__":
    main()
