src/sepsis_mcp/appendix_score_diagnostics.py

from __future__ import annotations

import argparse
import json
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from scipy.stats import ks_2samp

from sepsis_mcp.conformal import binary_nonconformity_scores
from sepsis_mcp.gossis import build_hospital_disjoint_split_with_selection, load_gossis_dataset
from sepsis_mcp.gossis_experiment import (
    GossisRunConfig,
    _build_structured_grouping,
    _encode_feature_splits,
    _normalize_json_value,
    _sanitize_encoded_features,
)
from sepsis_mcp.mimic4_experiment import (
    DEFAULT_MIN_UNIT_SIZE,
    _leave_one_out_assignments,
    _limit_frame_rows,
    _selection_features,
    load_mimic4_dataset,
)
from sepsis_mcp.modeling import ProbabilityEstimator


@dataclass(frozen=True)
class GossisScoreDiagnosticsConfig:
    data_root: Path
    min_hospital_admissions: int = 500
    selection_fraction: float = 0.1
    train_fraction: float = 0.6
    calibration_fraction: float = 0.1
    random_state_grid: tuple[int, ...] = (0,)
    model_type_grid: tuple[str, ...] = ("xgboost",)
    alpha: float = 0.1
    missingness_grouping_strategy: str = "coverage_gap_variable"


@dataclass(frozen=True)
class MimicScoreDiagnosticsConfig:
    csv_path: Path
    min_unit_size: int = DEFAULT_MIN_UNIT_SIZE
    rotations: int = 1
    max_assignments: int | None = 1
    model_type_grid: tuple[str, ...] = ("xgboost",)
    alpha: float = 0.1
    min_variable_missing_fraction: float = 0.1
    max_variable_missing_fraction: float = 0.9
    min_selection_group_rows: int = 100
    max_selection_rows: int | None = None
    max_calibration_rows: int | None = None
    max_test_rows: int | None = None


def _parse_int_grid(value: str) -> tuple[int, ...]:
    return tuple(int(item.strip()) for item in value.split(",") if item.strip())


def _parse_str_grid(value: str) -> tuple[str, ...]:
    return tuple(item.strip() for item in value.split(",") if item.strip())


def _score_records(
    *,
    dataset: str,
    model_type: str,
    seed: int,
    split_kind: str,
    scores: np.ndarray,
    groups: pd.DataFrame,
) -> pd.DataFrame:
    rows: list[dict[str, Any]] = []
    group_labels = groups["group_label"].astype(str).to_numpy()
    group_ids = groups["group"].astype(int).to_numpy()
    for score, group_id, group_label in zip(np.asarray(scores, dtype=float), group_ids, group_labels, strict=True):
        rows.append(
            {
                "dataset": dataset,
                "model_type": model_type,
                "seed": int(seed),
                "split_kind": split_kind,
                "group": int(group_id),
                "group_label": str(group_label),
                "score": float(score),
            }
        )
    return pd.DataFrame(rows)


def summarize_score_distribution_shifts(score_records: pd.DataFrame) -> pd.DataFrame:
    rows: list[dict[str, Any]] = []
    group_keys = ["dataset", "model_type", "seed"]
    for key, frame in score_records.groupby(group_keys, sort=True):
        context = dict(zip(group_keys, key if isinstance(key, tuple) else (key,), strict=False))
        calibration = frame[frame["split_kind"] == "calibration"].copy()
        pooled_scores = calibration["score"].to_numpy(dtype=float)
        for group_id, group_frame in calibration.groupby("group", sort=True):
            group_scores = group_frame["score"].to_numpy(dtype=float)
            pooled_ks = float(ks_2samp(group_scores, pooled_scores, method="auto").statistic)
            test_group = frame[(frame["split_kind"] == "test") & (frame["group"] == group_id)]
            calibration_vs_test_ks = float("nan")
            if not test_group.empty:
                calibration_vs_test_ks = float(
                    ks_2samp(group_scores, test_group["score"].to_numpy(dtype=float), method="auto").statistic
                )
            rows.append(
                {
                    **context,
                    "group": int(group_id),
                    "group_label": str(group_frame["group_label"].iloc[0]),
                    "pooled_ks": pooled_ks,
                    "calibration_vs_test_ks": calibration_vs_test_ks,
                    "n_calibration": int(len(group_frame)),
                    "n_test": int(len(test_group)),
                }
            )
    return pd.DataFrame(rows)


def _ecdf(values: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    ordered = np.sort(np.asarray(values, dtype=float))
    if ordered.size == 0:
        return np.array([]), np.array([])
    return ordered, np.arange(1, ordered.size + 1, dtype=float) / float(ordered.size)


def render_score_distribution_panels(score_records: pd.DataFrame, output_path: Path) -> None:
    panels = list(score_records.groupby(["dataset", "model_type"], sort=True))
    n_panels = max(len(panels), 1)
    fig, axes = plt.subplots(1, n_panels, figsize=(4.8 * n_panels, 3.4), squeeze=False)
    colors = ["#2171b5", "#e6550d", "#31a354", "#756bb1", "#636363"]
    for axis, ((dataset, model_type), frame) in zip(axes[0], panels, strict=False):
        calibration = frame[frame["split_kind"] == "calibration"]
        test = frame[frame["split_kind"] == "test"]
        pooled_x, pooled_y = _ecdf(calibration["score"].to_numpy(dtype=float))
        axis.plot(pooled_x, pooled_y, color="#111111", linewidth=2, label="Calibration pooled")
        for idx, (group_id, group_frame) in enumerate(calibration.groupby("group", sort=True)):
            color = colors[idx % len(colors)]
            label = f"Cal {group_frame['group_label'].iloc[0]}"
            x, y = _ecdf(group_frame["score"].to_numpy(dtype=float))
            axis.plot(x, y, color=color, linewidth=1.8, label=label)
            test_group = test[test["group"] == group_id]
            if not test_group.empty:
                x_test, y_test = _ecdf(test_group["score"].to_numpy(dtype=float))
                axis.plot(x_test, y_test, color=color, linewidth=1.4, linestyle="--", label=f"Test {group_frame['group_label'].iloc[0]}")
        axis.set_title(f"{dataset.upper()} / {model_type}")
        axis.set_xlabel("Nonconformity score")
        axis.set_ylabel("ECDF")
        axis.legend(frameon=False, fontsize=8)
    fig.tight_layout()
    output_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output_path, dpi=200)
    plt.close(fig)


def run_gossis_score_diagnostics(config: GossisScoreDiagnosticsConfig) -> pd.DataFrame:
    dataset = load_gossis_dataset(config.data_root, min_hospital_admissions=config.min_hospital_admissions)
    frames: list[pd.DataFrame] = []
    for seed in config.random_state_grid:
        split = build_hospital_disjoint_split_with_selection(
            dataset.frame,
            train_fraction=config.train_fraction,
            selection_fraction=config.selection_fraction,
            calibration_fraction=config.calibration_fraction,
            random_state=seed,
        )
        train_features, calibration_features, test_features = _encode_feature_splits(
            split.train_frame,
            split.calibration_frame,
            split.test_frame,
            dataset.feature_columns,
        )
        selection_features = (
            _sanitize_encoded_features(pd.get_dummies(split.selection_frame[dataset.feature_columns], dummy_na=True))
            .reindex(columns=train_features.columns, fill_value=0.0)
            .astype("float32")
        )
        for model_type in config.model_type_grid:
            estimator = ProbabilityEstimator(random_state=seed, model_type=model_type).fit(
                train_features,
                split.train_frame["label"],
            )
            selection_probabilities = estimator.predict_positive_proba(selection_features)
            calibration_probabilities = estimator.predict_positive_proba(calibration_features)
            test_probabilities = estimator.predict_positive_proba(test_features)
            grouping = _build_structured_grouping(
                config=GossisRunConfig(
                    data_root=config.data_root,
                    alpha=config.alpha,
                    train_fraction=config.train_fraction,
                    selection_fraction=config.selection_fraction,
                    calibration_fraction=config.calibration_fraction,
                    min_hospital_admissions=config.min_hospital_admissions,
                    random_state=seed,
                    model_type=model_type,
                    missingness_grouping_strategy=config.missingness_grouping_strategy,
                ),
                feature_columns=dataset.feature_columns,
                selection_frame=split.selection_frame,
                calibration_frame=split.calibration_frame,
                test_frame=split.test_frame,
                selection_probabilities=selection_probabilities,
                calibration_probabilities=calibration_probabilities,
                test_probabilities=test_probabilities,
            )
            calibration_scores = binary_nonconformity_scores(
                split.calibration_frame["label"].to_numpy(dtype=int),
                calibration_probabilities.to_numpy(dtype=float),
            )
            test_scores = binary_nonconformity_scores(
                split.test_frame["label"].to_numpy(dtype=int),
                test_probabilities.to_numpy(dtype=float),
            )
            frames.append(
                _score_records(
                    dataset="gossis",
                    model_type=model_type,
                    seed=seed,
                    split_kind="calibration",
                    scores=calibration_scores,
                    groups=grouping.calibration_groups,
                )
            )
            frames.append(
                _score_records(
                    dataset="gossis",
                    model_type=model_type,
                    seed=seed,
                    split_kind="test",
                    scores=test_scores,
                    groups=grouping.test_groups,
                )
            )
    return pd.concat(frames, ignore_index=True) if frames else pd.DataFrame()


def run_mimic_score_diagnostics(config: MimicScoreDiagnosticsConfig) -> pd.DataFrame:
    dataset = load_mimic4_dataset(config.csv_path, min_unit_size=config.min_unit_size)
    assignments = _leave_one_out_assignments(dataset.units, rotations=config.rotations)
    if config.max_assignments is not None:
        assignments = assignments[: config.max_assignments]
    frames: list[pd.DataFrame] = []
    for assignment in assignments:
        seed = int(assignment["rotation"])
        frame = dataset.frame
        train_frame = frame[frame["hospital_id"].isin(assignment["train_units"])].copy()
        selection_frame = _limit_frame_rows(
            frame[frame["hospital_id"].isin(assignment["selection_units"])].copy(),
            config.max_selection_rows,
            random_state=seed + 201,
        )
        calibration_frame = _limit_frame_rows(
            frame[frame["hospital_id"] == assignment["calibration_unit"]].copy(),
            config.max_calibration_rows,
            random_state=seed + 203,
        )
        test_frame = _limit_frame_rows(
            frame[frame["hospital_id"] == assignment["test_unit"]].copy(),
            config.max_test_rows,
            random_state=seed + 207,
        )
        train_features, calibration_features, test_features = _encode_feature_splits(
            train_frame,
            calibration_frame,
            test_frame,
            dataset.feature_columns,
        )
        selection_features = _selection_features(
            selection_frame,
            train_columns=train_features.columns,
            feature_columns=dataset.feature_columns,
        )
        for model_type in config.model_type_grid:
            estimator = ProbabilityEstimator(random_state=seed, model_type=model_type).fit(
                train_features,
                train_frame["label"],
            )
            selection_probabilities = estimator.predict_positive_proba(selection_features)
            calibration_probabilities = estimator.predict_positive_proba(calibration_features)
            test_probabilities = estimator.predict_positive_proba(test_features)
            grouping = _build_structured_grouping(
                config=GossisRunConfig(
                    data_root=Path("."),
                    alpha=config.alpha,
                    selection_fraction=0.1,
                    model_type=model_type,
                    random_state=seed,
                    missingness_grouping_strategy="coverage_gap_variable",
                    min_variable_missing_fraction=config.min_variable_missing_fraction,
                    max_variable_missing_fraction=config.max_variable_missing_fraction,
                    min_selection_group_rows=config.min_selection_group_rows,
                ),
                feature_columns=dataset.feature_columns,
                selection_frame=selection_frame,
                calibration_frame=calibration_frame,
                test_frame=test_frame,
                selection_probabilities=selection_probabilities,
                calibration_probabilities=calibration_probabilities,
                test_probabilities=test_probabilities,
            )
            calibration_scores = binary_nonconformity_scores(
                calibration_frame["label"].to_numpy(dtype=int),
                calibration_probabilities.to_numpy(dtype=float),
            )
            test_scores = binary_nonconformity_scores(
                test_frame["label"].to_numpy(dtype=int),
                test_probabilities.to_numpy(dtype=float),
            )
            frames.append(
                _score_records(
                    dataset="mimic4",
                    model_type=model_type,
                    seed=seed,
                    split_kind="calibration",
                    scores=calibration_scores,
                    groups=grouping.calibration_groups,
                )
            )
            frames.append(
                _score_records(
                    dataset="mimic4",
                    model_type=model_type,
                    seed=seed,
                    split_kind="test",
                    scores=test_scores,
                    groups=grouping.test_groups,
                )
            )
    return pd.concat(frames, ignore_index=True) if frames else pd.DataFrame()


def run_score_diagnostics(
    *,
    output_dir: Path,
    gossis_config: GossisScoreDiagnosticsConfig | None = None,
    mimic_config: MimicScoreDiagnosticsConfig | None = None,
) -> dict[str, Path]:
    output_dir.mkdir(parents=True, exist_ok=True)
    frames: list[pd.DataFrame] = []
    config_payload: dict[str, Any] = {}
    if gossis_config is not None:
        frames.append(run_gossis_score_diagnostics(gossis_config))
        config_payload["gossis"] = asdict(gossis_config)
    if mimic_config is not None:
        frames.append(run_mimic_score_diagnostics(mimic_config))
        config_payload["mimic4"] = asdict(mimic_config)
    score_records = pd.concat([frame for frame in frames if not frame.empty], ignore_index=True) if frames else pd.DataFrame()
    score_summary = summarize_score_distribution_shifts(score_records) if not score_records.empty else pd.DataFrame()

    records_path = output_dir / "score_records.csv"
    summary_path = output_dir / "score_summary.csv"
    figure_path = output_dir / "score_ecdf_panels.png"
    config_path = output_dir / "score_diagnostics_config.json"

    score_records.to_csv(records_path, index=False)
    score_summary.to_csv(summary_path, index=False)
    if not score_records.empty:
        render_score_distribution_panels(score_records, figure_path)
    config_path.write_text(json.dumps(_normalize_json_value(config_payload), indent=2, sort_keys=True), encoding="utf-8")

    return {
        "score_records": records_path,
        "score_summary": summary_path,
        "score_figure": figure_path,
        "config": config_path,
    }


def main(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(description="Appendix score distribution diagnostics for GOSSIS and MIMIC-IV")
    parser.add_argument("--gossis-data-root", type=Path, default=None)
    parser.add_argument("--mimic-csv-path", type=Path, default=None)
    parser.add_argument("--min-hospital-admissions", type=int, default=500)
    parser.add_argument("--min-unit-size", type=int, default=DEFAULT_MIN_UNIT_SIZE)
    parser.add_argument("--random-state-grid", type=_parse_int_grid, default=(0,))
    parser.add_argument("--model-type-grid", type=_parse_str_grid, default=("xgboost",))
    parser.add_argument("--rotations", type=int, default=1)
    parser.add_argument("--max-assignments", type=int, default=1)
    parser.add_argument("--output-dir", type=Path, required=True)
    args = parser.parse_args(argv)

    gossis_config = None
    if args.gossis_data_root is not None:
        gossis_config = GossisScoreDiagnosticsConfig(
            data_root=args.gossis_data_root,
            min_hospital_admissions=args.min_hospital_admissions,
            random_state_grid=tuple(args.random_state_grid),
            model_type_grid=tuple(args.model_type_grid),
        )
    mimic_config = None
    if args.mimic_csv_path is not None:
        mimic_config = MimicScoreDiagnosticsConfig(
            csv_path=args.mimic_csv_path,
            min_unit_size=args.min_unit_size,
            rotations=args.rotations,
            max_assignments=args.max_assignments,
            model_type_grid=tuple(args.model_type_grid),
        )
    run_score_diagnostics(output_dir=args.output_dir, gossis_config=gossis_config, mimic_config=mimic_config)


if __name__ == "__main__":
    main()