src/sepsis_mcp/appendix_negative_case_analysis.py

from __future__ import annotations

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

import pandas as pd

from sepsis_mcp.simulation import SimulationConfig, run_sim3_sweep


@dataclass
class NegativeCaseConfig:
    output_dir: Path
    delta_m_grid: tuple[float, ...] = (0.0, 0.05, 0.1)
    delta_x: float = 0.0
    gamma: float = 0.0
    seeds: int = 10
    model_type: str = "logistic_regression"
    n_train: int = 5000
    n_cal: int = 2000
    n_test: int = 3000
    d: int = 20
    alpha: float = 0.1
    base_missing_rate: float = 0.05


def build_negative_case_summary(summary_frame: pd.DataFrame) -> pd.DataFrame:
    if summary_frame.empty:
        return pd.DataFrame()
    focus = summary_frame[
        summary_frame["mechanism"].eq("mcar")
        & summary_frame["method"].isin(["standard", "mondrian_tilted"])
    ].copy()
    if focus.empty:
        return focus
    pivot = focus.pivot_table(
        index=["delta_m", "mechanism"],
        columns="method",
        values=[
            "max_group_coverage_gap_mean",
            "empirical_coverage_mean",
            "average_set_size_mean",
        ],
    )
    pivot.columns = [f"{metric}_{method}" for metric, method in pivot.columns]
    result = pivot.reset_index()
    result["gap_advantage_vs_standard"] = (
        result["max_group_coverage_gap_mean_standard"]
        - result["max_group_coverage_gap_mean_mondrian_tilted"]
    )
    result["coverage_difference_vs_standard"] = (
        result["empirical_coverage_mean_mondrian_tilted"]
        - result["empirical_coverage_mean_standard"]
    )
    result["set_size_difference_vs_standard"] = (
        result["average_set_size_mean_mondrian_tilted"]
        - result["average_set_size_mean_standard"]
    )
    result["is_neutral_or_worse"] = result["gap_advantage_vs_standard"] <= 0.0
    return result.sort_values(["delta_m", "mechanism"]).reset_index(drop=True)


def run_negative_case_analysis(config: NegativeCaseConfig) -> dict[str, Path]:
    config.output_dir.mkdir(parents=True, exist_ok=True)
    summary_rows: list[pd.DataFrame] = []
    for delta_m in config.delta_m_grid:
        sweep_dir = config.output_dir / "runs" / f"delta_m_{str(float(delta_m)).replace('.', 'p')}"
        _, summary = run_sim3_sweep(
            SimulationConfig(
                n_train=config.n_train,
                n_cal=config.n_cal,
                n_test=config.n_test,
                d=config.d,
                alpha=config.alpha,
                base_missing_rate=config.base_missing_rate,
                model_type=config.model_type,
            ),
            delta_m=float(delta_m),
            delta_x=float(config.delta_x),
            gamma=float(config.gamma),
            seeds=int(config.seeds),
            output_dir=sweep_dir,
        )
        summary = summary.copy()
        summary["delta_m"] = float(delta_m)
        summary_rows.append(summary)

    combined_summary = pd.concat(summary_rows, ignore_index=True)
    negative_case_summary = build_negative_case_summary(combined_summary)
    combined_path = config.output_dir / "simulation_negative_case_summary.csv"
    focus_path = config.output_dir / "negative_case_focus.csv"
    config_path = config.output_dir / "config.json"
    manifest_path = config.output_dir / "manifest.json"
    combined_summary.to_csv(combined_path, index=False)
    negative_case_summary.to_csv(focus_path, index=False)
    config_path.write_text(json.dumps(asdict(config), indent=2, default=str), encoding="utf-8")
    manifest_path.write_text(
        json.dumps(
            {
                "simulation_summary": str(combined_path),
                "negative_case_focus": str(focus_path),
                "config": str(config_path),
            },
            indent=2,
            sort_keys=True,
        ),
        encoding="utf-8",
    )
    return {
        "simulation_summary": combined_path,
        "negative_case_focus": focus_path,
        "config": config_path,
        "manifest": manifest_path,
    }


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


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="sepsis-mcp-appendix-negative-case-analysis")
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--delta-m-grid", type=_parse_float_grid, default="0.0,0.05,0.1")
    parser.add_argument("--delta-x", type=float, default=0.0)
    parser.add_argument("--gamma", type=float, default=0.0)
    parser.add_argument("--seeds", type=int, default=10)
    parser.add_argument("--model-type", choices={"logistic_regression", "xgboost"}, default="logistic_regression")
    parser.add_argument("--n-train", type=int, default=5000)
    parser.add_argument("--n-cal", type=int, default=2000)
    parser.add_argument("--n-test", type=int, default=3000)
    parser.add_argument("--d", type=int, default=20)
    parser.add_argument("--alpha", type=float, default=0.1)
    parser.add_argument("--base-missing-rate", type=float, default=0.05)
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    run_negative_case_analysis(
        NegativeCaseConfig(
            output_dir=args.output_dir,
            delta_m_grid=args.delta_m_grid,
            delta_x=args.delta_x,
            gamma=args.gamma,
            seeds=args.seeds,
            model_type=args.model_type,
            n_train=args.n_train,
            n_cal=args.n_cal,
            n_test=args.n_test,
            d=args.d,
            alpha=args.alpha,
            base_missing_rate=args.base_missing_rate,
        )
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())