#!/usr/bin/env python3
"""Analyze annual CDC/NVSS natality aggregates for a modified Robson study.

Inputs are aggregate JSON files produced by ``aggregate_natality.py``. The
analysis never reads or writes person-level birth records.
"""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import norm


GROUP_LABELS = {
    "1": "Nulliparous, term, singleton, cephalic, spontaneous labor",
    "2": "Nulliparous, term, singleton, cephalic, induced/prelabor cesarean",
    "3": "Multiparous without previous cesarean, term, singleton, cephalic, spontaneous labor",
    "4": "Multiparous without previous cesarean, term, singleton, cephalic, induced/prelabor cesarean",
    "5": "Previous cesarean, term, singleton, cephalic",
    "6": "Nulliparous singleton breech",
    "7": "Multiparous singleton breech",
    "8": "Multiple pregnancy",
    "9": "Singleton with other/non-cephalic, non-breech presentation",
    "10": "Singleton cephalic, preterm",
}

GROUP_SHORT = {
    "1": "Group 1",
    "2": "Group 2",
    "3": "Group 3",
    "4": "Group 4",
    "5": "Group 5",
    "6": "Group 6",
    "7": "Group 7",
    "8": "Group 8",
    "9": "Group 9*",
    "10": "Group 10",
}

PALETTE = {
    "navy": "#1F4E79",
    "gold": "#D5963B",
    "gray": "#64748B",
    "light_blue": "#B9C9D8",
    "light_gray": "#E5E7EB",
    "ink": "#1F2937",
}


def wilson_interval(successes: int, total: int, alpha: float = 0.05) -> tuple[float, float]:
    if total <= 0:
        return float("nan"), float("nan")
    z = norm.ppf(1 - alpha / 2)
    p = successes / total
    denominator = 1 + z * z / total
    center = (p + z * z / (2 * total)) / denominator
    half = (
        z
        * math.sqrt((p * (1 - p) / total) + z * z / (4 * total * total))
        / denominator
    )
    return center - half, center + half


def poisson_loglinear(
    years: np.ndarray,
    events: np.ndarray,
    exposure: np.ndarray,
    *,
    label: str,
) -> dict:
    """Fit a quasi-Poisson log-linear rate model by Newton scoring."""

    centered = years - years.min()
    design = np.column_stack([np.ones(len(years)), centered])
    beta = np.array([math.log(events.sum() / exposure.sum()), 0.0])
    offset = np.log(exposure)
    iterations = 0
    for iterations in range(1, 101):
        mu = np.exp(np.clip(offset + design @ beta, -50, 50))
        information = design.T @ (mu[:, None] * design)
        score = design.T @ (events - mu)
        step = np.linalg.solve(information, score)
        beta_next = beta + step
        if np.max(np.abs(step)) < 1e-12:
            beta = beta_next
            break
        beta = beta_next

    mu = np.exp(np.clip(offset + design @ beta, -50, 50))
    information = design.T @ (mu[:, None] * design)
    residual_df = len(events) - design.shape[1]
    pearson = float(np.sum((events - mu) ** 2 / mu))
    dispersion_raw = pearson / residual_df if residual_df > 0 else 1.0
    dispersion = max(1.0, dispersion_raw)
    covariance = np.linalg.inv(information) * dispersion
    se = float(math.sqrt(covariance[1, 1]))
    slope = float(beta[1])
    z_value = slope / se
    p_value = float(2 * norm.sf(abs(z_value)))
    annual_percent_change = (math.exp(slope) - 1) * 100
    lower = (math.exp(slope - 1.96 * se) - 1) * 100
    upper = (math.exp(slope + 1.96 * se) - 1) * 100
    return {
        "model": label,
        "first_year": int(years.min()),
        "last_year": int(years.max()),
        "n_time_points": int(len(years)),
        "annual_percent_change": annual_percent_change,
        "ci_lower": lower,
        "ci_upper": upper,
        "p_value": p_value,
        "dispersion_raw": dispersion_raw,
        "dispersion_used": dispersion,
        "iterations": iterations,
    }


def set_plot_style() -> None:
    plt.rcParams.update(
        {
            "font.family": "DejaVu Sans",
            "font.size": 9,
            "axes.titlesize": 12,
            "axes.labelsize": 10,
            "axes.edgecolor": "#94A3B8",
            "axes.linewidth": 0.8,
            "xtick.color": PALETTE["ink"],
            "ytick.color": PALETTE["ink"],
            "text.color": PALETTE["ink"],
            "figure.facecolor": "white",
            "axes.facecolor": "white",
            "savefig.facecolor": "white",
            "savefig.bbox": "tight",
        }
    )


def save_figure(fig: plt.Figure, stem: Path) -> None:
    stem.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(stem.with_suffix(".png"), dpi=300)
    fig.savefig(stem.with_suffix(".tiff"), dpi=300, pil_kwargs={"compression": "tiff_lzw"})
    plt.close(fig)


def plot_flow(latest: dict, output_dir: Path) -> None:
    flow = latest["flow"]
    stages = [
        ("Raw 2024 live-birth records", flow["raw_records"]),
        ("U.S. residents", flow["us_resident"]),
        ("Hospital births", flow["hospital_us_resident"]),
        ("Known delivery method", flow["known_delivery"]),
        ("Modified Robson classified", flow["robson_classified"]),
    ]
    exclusions = [
        flow["excluded_non_us_resident"],
        flow["excluded_nonhospital_or_unknown"],
        flow["excluded_delivery_unknown"],
        flow["excluded_unclassifiable_robson"],
    ]
    fig, ax = plt.subplots(figsize=(7.2, 8.0))
    ax.set_xlim(0, 10)
    ax.set_ylim(0, 11)
    ax.axis("off")
    y_positions = [9.6, 7.7, 5.8, 3.9, 2.0]
    for idx, ((label, count), y) in enumerate(zip(stages, y_positions)):
        face = PALETTE["navy"] if idx in {0, 4} else "#EEF3F8"
        color = "white" if idx in {0, 4} else PALETTE["ink"]
        ax.add_patch(
            plt.Rectangle(
                (1.5, y - 0.55),
                7.0,
                1.1,
                facecolor=face,
                edgecolor=PALETTE["navy"],
                linewidth=1.2,
                zorder=2,
            )
        )
        ax.text(
            5,
            y,
            f"{label}\n{count:,}",
            ha="center",
            va="center",
            color=color,
            fontsize=10,
            fontweight="bold" if idx in {0, 4} else "normal",
            zorder=3,
        )
        if idx < len(stages) - 1:
            next_y = y_positions[idx + 1]
            ax.annotate(
                "",
                xy=(5, next_y + 0.6),
                xytext=(5, y - 0.6),
                arrowprops=dict(arrowstyle="-|>", color=PALETTE["gray"], lw=1.2),
            )
            ax.text(
                8.8,
                (y + next_y) / 2,
                f"Excluded\n{exclusions[idx]:,}",
                ha="left",
                va="center",
                fontsize=8.5,
                color=PALETTE["gray"],
            )
    ax.text(
        0,
        0.45,
        "Source: CDC/NCHS National Vital Statistics System natality public-use file, 2024.",
        fontsize=8,
        color=PALETTE["gray"],
    )
    save_figure(fig, output_dir / "figure_1_study_flow")


def plot_overall(annual: list[dict], output_dir: Path) -> None:
    years = np.array([row["year"] for row in annual])
    rates = np.array([row["cesarean_rate"] for row in annual]) * 100
    rates_ex9 = np.array([row["cesarean_rate_excluding_group9"] for row in annual]) * 100
    fig, ax = plt.subplots(figsize=(8.0, 4.9))
    ax.plot(
        years,
        rates,
        color=PALETTE["navy"],
        marker="o",
        lw=2.3,
        label="Primary modified Robson population",
    )
    ax.plot(
        years,
        rates_ex9,
        color=PALETTE["gold"],
        marker="s",
        ms=4,
        lw=1.8,
        linestyle="--",
        label="Sensitivity: excluding modified group 9",
    )
    ax.annotate(
        f"{rates[0]:.2f}%",
        (years[0], rates[0]),
        xytext=(5, 11),
        textcoords="offset points",
        fontsize=8.5,
        color=PALETTE["navy"],
    )
    ax.annotate(
        f"{rates[-1]:.2f}%",
        (years[-1], rates[-1]),
        xytext=(-34, 11),
        textcoords="offset points",
        fontsize=8.5,
        color=PALETTE["navy"],
    )
    ax.set_ylabel("Cesarean deliveries (%)")
    ax.set_xlabel("Year")
    ax.set_xticks(years)
    ax.set_ylim(min(rates_ex9.min(), rates.min()) - 0.25, max(rates.max(), rates_ex9.max()) + 0.3)
    ax.grid(axis="y", color=PALETTE["light_gray"], lw=0.8)
    ax.spines[["top", "right"]].set_visible(False)
    ax.legend(frameon=False, loc="upper left", fontsize=8.5)
    ax.text(
        0,
        -0.23,
        "Population: U.S.-resident hospital live births with known delivery method and sufficient variables for modified Robson classification.",
        transform=ax.transAxes,
        fontsize=7.8,
        color=PALETTE["gray"],
    )
    save_figure(fig, output_dir / "figure_2_overall_trend")


def plot_2024_contribution(rows_2024: list[dict], output_dir: Path) -> None:
    ordered = sorted(rows_2024, key=lambda row: row["relative_contribution"], reverse=True)
    labels = [GROUP_SHORT[row["group"]] for row in ordered][::-1]
    values = [row["relative_contribution"] * 100 for row in ordered][::-1]
    fig, ax = plt.subplots(figsize=(8.0, 5.7))
    bars = ax.barh(labels, values, color=PALETTE["navy"], height=0.68)
    for bar, value in zip(bars, values):
        ax.text(
            value + 0.35,
            bar.get_y() + bar.get_height() / 2,
            f"{value:.1f}%",
            va="center",
            fontsize=8.5,
        )
    ax.set_xlabel("Share of all cesarean deliveries (%)")
    ax.set_xlim(0, max(values) + 5.0)
    ax.grid(axis="x", color=PALETTE["light_gray"], lw=0.8)
    ax.spines[["top", "right", "left"]].set_visible(False)
    ax.tick_params(axis="y", length=0)
    ax.text(
        0,
        -0.16,
        "*Modified group 9 is based on the NVSS category “other presentation” and does not map exactly to transverse/oblique lie.",
        transform=ax.transAxes,
        fontsize=7.8,
        color=PALETTE["gray"],
    )
    save_figure(fig, output_dir / "figure_3_2024_contribution")


def plot_decomposition(decomposition: list[dict], output_dir: Path) -> None:
    rows = sorted(decomposition, key=lambda row: int(row["group"]))
    labels = [GROUP_SHORT[row["group"]] for row in rows]
    composition = np.array([row["composition_effect_pp"] for row in rows])
    rate = np.array([row["rate_effect_pp"] for row in rows])
    y = np.arange(len(rows))
    fig, ax = plt.subplots(figsize=(8.2, 6.0))
    ax.barh(
        y - 0.18,
        composition,
        height=0.32,
        color=PALETTE["navy"],
        label="Population-share effect",
    )
    ax.barh(
        y + 0.18,
        rate,
        height=0.32,
        color=PALETTE["gold"],
        label="Within-group rate effect",
    )
    ax.axvline(0, color=PALETTE["ink"], lw=0.8)
    ax.set_yticks(y, labels)
    ax.invert_yaxis()
    ax.set_xlabel("Contribution to absolute change in overall cesarean rate (percentage points)")
    ax.grid(axis="x", color=PALETTE["light_gray"], lw=0.8)
    ax.spines[["top", "right", "left"]].set_visible(False)
    ax.tick_params(axis="y", length=0)
    ax.legend(frameon=False, loc="lower right", fontsize=8.5)
    ax.text(
        0,
        -0.14,
        "Positive values increased the overall rate; negative values offset the increase. Effects sum to the observed absolute change.",
        transform=ax.transAxes,
        fontsize=7.8,
        color=PALETTE["gray"],
    )
    save_figure(fig, output_dir / "figure_4_decomposition")


def plot_group_trends(group_rows: list[dict], output_dir: Path) -> None:
    fig, axes = plt.subplots(5, 2, figsize=(9.0, 12.0), sharex=True)
    axes = axes.ravel()
    for ax, group in zip(axes, [str(i) for i in range(1, 11)]):
        rows = sorted(
            [row for row in group_rows if row["group"] == group],
            key=lambda row: row["year"],
        )
        years = [row["year"] for row in rows]
        rates = [row["cesarean_rate"] * 100 for row in rows]
        ax.plot(years, rates, color=PALETTE["navy"], marker="o", ms=3, lw=1.7)
        ax.set_title(GROUP_SHORT[group], loc="left", fontsize=10, fontweight="bold")
        ax.grid(axis="y", color=PALETTE["light_gray"], lw=0.7)
        ax.spines[["top", "right"]].set_visible(False)
        ax.set_xticks([2016, 2018, 2020, 2022, 2024])
        ax.tick_params(labelsize=8)
        ax.set_ylabel("%", rotation=0, labelpad=8, fontsize=8)
    fig.text(
        0.08,
        0.005,
        "*Modified group 9 uses the NVSS category “other presentation”; interpret its level and trend cautiously.",
        fontsize=8,
        color=PALETTE["gray"],
    )
    fig.tight_layout(rect=(0, 0.025, 1, 1))
    save_figure(fig, output_dir / "supplementary_figure_1_group_trends")


def analyze(input_dir: Path, output_dir: Path) -> dict:
    records = []
    for path in sorted(input_dir.glob("aggregate_20??.json")):
        with path.open(encoding="utf-8") as handle:
            records.append(json.load(handle))
    years_found = [record["year"] for record in records]
    if years_found != list(range(2016, 2025)):
        raise RuntimeError(f"Expected 2016-2024 aggregates, found {years_found}")

    annual = []
    group_rows = []
    subgroup_rows = []
    modifier_rows = []
    quality_rows = []
    provenance_rows = []

    for record in records:
        year = record["year"]
        flow = record["flow"]
        group_n = sum(item["n"] for item in record["groups"].values())
        group_cs = sum(item["cesarean"] for item in record["groups"].values())
        if group_n != flow["robson_classified"]:
            raise AssertionError(f"{year}: group total does not reconcile with classified flow")
        if group_cs <= 0 or group_cs >= group_n:
            raise AssertionError(f"{year}: implausible cesarean count")
        group9 = record["groups"]["9"]
        rate_low, rate_high = wilson_interval(group_cs, group_n)
        annual.append(
            {
                "year": year,
                "raw_records": flow["raw_records"],
                "hospital_us_resident": flow["hospital_us_resident"],
                "known_delivery": flow["known_delivery"],
                "classified": group_n,
                "cesarean": group_cs,
                "cesarean_rate": group_cs / group_n,
                "cesarean_rate_ci_lower": rate_low,
                "cesarean_rate_ci_upper": rate_high,
                "classification_completeness": group_n / flow["known_delivery"],
                "unclassifiable": flow["excluded_unclassifiable_robson"],
                "classified_excluding_group9": group_n - group9["n"],
                "cesarean_excluding_group9": group_cs - group9["cesarean"],
                "cesarean_rate_excluding_group9": (
                    (group_cs - group9["cesarean"]) / (group_n - group9["n"])
                ),
            }
        )

        for group in [str(i) for i in range(1, 11)]:
            values = record["groups"][group]
            low, high = wilson_interval(values["cesarean"], values["n"])
            group_rows.append(
                {
                    "year": year,
                    "group": group,
                    "definition": GROUP_LABELS[group],
                    "n": values["n"],
                    "population_share": values["n"] / group_n,
                    "cesarean": values["cesarean"],
                    "cesarean_rate": values["cesarean"] / values["n"],
                    "cesarean_rate_ci_lower": low,
                    "cesarean_rate_ci_upper": high,
                    "relative_contribution": values["cesarean"] / group_cs,
                    "apgar_0_3": values.get("apgar5_0_3", 0),
                    "apgar_known": values.get("apgar5_known", 0),
                    "apgar_0_3_rate": (
                        values.get("apgar5_0_3", 0) / values.get("apgar5_known", 1)
                    ),
                    "nicu": values.get("nicu", 0),
                    "nicu_known": values.get("nicu_known", 0),
                    "nicu_rate": values.get("nicu", 0) / values.get("nicu_known", 1),
                    "ventilation_6h": values.get("vent6", 0),
                    "ventilation_6h_known": values.get("vent6_known", 0),
                    "ventilation_6h_rate": (
                        values.get("vent6", 0) / values.get("vent6_known", 1)
                    ),
                }
            )

        for subgroup in ["2a", "2b", "4a", "4b"]:
            values = record["group_subgroups"][subgroup]
            subgroup_rows.append(
                {
                    "year": year,
                    "subgroup": subgroup,
                    "n": values["n"],
                    "population_share": values["n"] / group_n,
                    "cesarean": values["cesarean"],
                    "cesarean_rate": values["cesarean"] / values["n"],
                    "relative_contribution": values["cesarean"] / group_cs,
                }
            )

        for key, values in sorted(record["modifiers"].items()):
            variable, category = key.split("|", 1)
            modifier_rows.append(
                {
                    "year": year,
                    "variable": variable,
                    "category": category,
                    "n": values["n"],
                    "population_share": values["n"] / group_n,
                    "cesarean": values["cesarean"],
                    "cesarean_rate": values["cesarean"] / values["n"],
                    "relative_contribution": values["cesarean"] / group_cs,
                }
            )

        missing = record["missing"]
        conflicts = record["conflicts"]
        quality_rows.append(
            {
                "year": year,
                "known_delivery": flow["known_delivery"],
                "classified": group_n,
                "classification_completeness": group_n / flow["known_delivery"],
                "missing_live_birth_order": missing.get("live_birth_order", 0),
                "missing_previous_cesareans": missing.get("previous_cesareans", 0),
                "missing_presentation": missing.get("presentation", 0),
                "missing_gestation": missing.get("gestation", 0),
                "missing_induction": missing.get("induction", 0),
                "missing_trial_labor_among_cesareans": missing.get(
                    "trial_labor_among_cesareans", 0
                ),
                "conflict_first_live_birth_with_previous_cesarean": conflicts.get(
                    "first_live_birth_with_previous_cesarean", 0
                ),
                "conflict_induction_yes_trial_no": conflicts.get(
                    "induction_yes_trial_no", 0
                ),
            }
        )
        provenance_rows.append(
            {
                "year": year,
                "source_zip": record["source_zip"],
                "source_member": record["source_member"],
                "source_zip_sha256": record["source_zip_sha256"],
                "source_member_crc": record["source_member_crc"],
                "source_member_uncompressed_bytes": record[
                    "source_member_uncompressed_bytes"
                ],
                "record_length_sample": ",".join(record["record_length_sample"].keys()),
            }
        )

    annual_by_year = {row["year"]: row for row in annual}
    group_lookup = {(row["year"], row["group"]): row for row in group_rows}
    decomposition = []
    for group in [str(i) for i in range(1, 11)]:
        row0 = group_lookup[(2016, group)]
        row1 = group_lookup[(2024, group)]
        p0, p1 = row0["population_share"], row1["population_share"]
        r0, r1 = row0["cesarean_rate"], row1["cesarean_rate"]
        composition_effect = (p1 - p0) * (r1 + r0) / 2
        rate_effect = (r1 - r0) * (p1 + p0) / 2
        decomposition.append(
            {
                "group": group,
                "definition": GROUP_LABELS[group],
                "population_share_2016": p0,
                "population_share_2024": p1,
                "population_share_change_pp": (p1 - p0) * 100,
                "cesarean_rate_2016": r0,
                "cesarean_rate_2024": r1,
                "cesarean_rate_change_pp": (r1 - r0) * 100,
                "composition_effect_pp": composition_effect * 100,
                "rate_effect_pp": rate_effect * 100,
                "total_effect_pp": (composition_effect + rate_effect) * 100,
            }
        )
    observed_change_pp = (
        annual_by_year[2024]["cesarean_rate"]
        - annual_by_year[2016]["cesarean_rate"]
    ) * 100
    decomposed_change_pp = sum(row["total_effect_pp"] for row in decomposition)
    if not math.isclose(observed_change_pp, decomposed_change_pp, abs_tol=1e-10):
        raise AssertionError("Kitagawa decomposition does not reconcile")

    years = np.array([row["year"] for row in annual], dtype=float)
    events = np.array([row["cesarean"] for row in annual], dtype=float)
    exposure = np.array([row["classified"] for row in annual], dtype=float)
    trends = [poisson_loglinear(years, events, exposure, label="Overall primary")]
    mask = ~np.isin(years, [2020, 2021])
    trends.append(
        poisson_loglinear(
            years[mask],
            events[mask],
            exposure[mask],
            label="Overall excluding 2020-2021",
        )
    )
    events_ex9 = np.array(
        [row["cesarean_excluding_group9"] for row in annual], dtype=float
    )
    exposure_ex9 = np.array(
        [row["classified_excluding_group9"] for row in annual], dtype=float
    )
    trends.append(
        poisson_loglinear(
            years,
            events_ex9,
            exposure_ex9,
            label="Overall excluding modified group 9",
        )
    )
    for group in [str(i) for i in range(1, 11)]:
        rows = [row for row in group_rows if row["group"] == group]
        trends.append(
            poisson_loglinear(
                np.array([row["year"] for row in rows], dtype=float),
                np.array([row["cesarean"] for row in rows], dtype=float),
                np.array([row["n"] for row in rows], dtype=float),
                label=f"Group {group}",
            )
        )

    pooled_groups = []
    total_classified = sum(row["classified"] for row in annual)
    total_cesarean = sum(row["cesarean"] for row in annual)
    for group in [str(i) for i in range(1, 11)]:
        rows = [row for row in group_rows if row["group"] == group]
        n = sum(row["n"] for row in rows)
        cesarean = sum(row["cesarean"] for row in rows)
        pooled_groups.append(
            {
                "group": group,
                "definition": GROUP_LABELS[group],
                "n": n,
                "population_share": n / total_classified,
                "cesarean": cesarean,
                "cesarean_rate": cesarean / n,
                "relative_contribution": cesarean / total_cesarean,
            }
        )

    result = {
        "metadata": {
            "analysis_title": "Modified Robson analysis of U.S. cesarean delivery, 2016-2024",
            "data_source": "CDC/NCHS National Vital Statistics System natality public-use files",
            "total_classified_births": total_classified,
            "total_cesarean_deliveries": total_cesarean,
            "overall_cesarean_rate": total_cesarean / total_classified,
            "observed_change_2016_2024_pp": observed_change_pp,
            "decomposed_change_2016_2024_pp": decomposed_change_pp,
            "robson_note": (
                "Modified group 9 uses the NVSS 'other presentation' category and "
                "does not map exactly to transverse/oblique lie."
            ),
        },
        "annual": annual,
        "group_year": group_rows,
        "subgroup_year": subgroup_rows,
        "modifier_year": modifier_rows,
        "quality": quality_rows,
        "provenance": provenance_rows,
        "decomposition": decomposition,
        "trend_models": trends,
        "pooled_groups": pooled_groups,
    }
    output_dir.mkdir(parents=True, exist_ok=True)
    with (output_dir / "analysis_results.json").open("w", encoding="utf-8") as handle:
        json.dump(result, handle, ensure_ascii=False, indent=2)
        handle.write("\n")

    set_plot_style()
    plot_flow(records[-1], output_dir)
    plot_overall(annual, output_dir)
    plot_2024_contribution(
        [row for row in group_rows if row["year"] == 2024], output_dir
    )
    plot_decomposition(decomposition, output_dir)
    plot_group_trends(group_rows, output_dir)
    return result


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input-dir", type=Path, default=Path("analysis"))
    parser.add_argument("--output-dir", type=Path, default=Path("outputs/analysis"))
    args = parser.parse_args()
    result = analyze(args.input_dir, args.output_dir)
    metadata = result["metadata"]
    print(
        json.dumps(
            {
                "total_classified_births": metadata["total_classified_births"],
                "total_cesarean_deliveries": metadata["total_cesarean_deliveries"],
                "observed_change_pp": metadata["observed_change_2016_2024_pp"],
            }
        )
    )


if __name__ == "__main__":
    main()
