from __future__ import annotations

from pathlib import Path

import pandas as pd

from tests.gossis_helpers import write_structured_synthetic_gossis_dataset
from tests.test_mimic4_experiment import _write_synthetic_mimic4_csv


def test_summarize_score_distribution_shifts_computes_groupwise_ks() -> None:
    from sepsis_mcp.appendix_score_diagnostics import summarize_score_distribution_shifts

    frame = pd.DataFrame(
        [
            {"dataset": "demo", "model_type": "xgboost", "seed": 0, "split_kind": "calibration", "group": 0, "group_label": "low", "score": 0.05},
            {"dataset": "demo", "model_type": "xgboost", "seed": 0, "split_kind": "calibration", "group": 0, "group_label": "low", "score": 0.10},
            {"dataset": "demo", "model_type": "xgboost", "seed": 0, "split_kind": "calibration", "group": 1, "group_label": "high", "score": 0.80},
            {"dataset": "demo", "model_type": "xgboost", "seed": 0, "split_kind": "calibration", "group": 1, "group_label": "high", "score": 0.85},
            {"dataset": "demo", "model_type": "xgboost", "seed": 0, "split_kind": "test", "group": 0, "group_label": "low", "score": 0.06},
            {"dataset": "demo", "model_type": "xgboost", "seed": 0, "split_kind": "test", "group": 1, "group_label": "high", "score": 0.95},
        ]
    )

    summary = summarize_score_distribution_shifts(frame)

    assert {"dataset", "model_type", "seed", "group", "group_label", "pooled_ks", "calibration_vs_test_ks"} <= set(
        summary.columns
    )
    assert set(summary["group_label"]) == {"low", "high"}
    assert (summary["pooled_ks"] > 0.0).all()


def test_run_score_diagnostics_writes_expected_outputs(tmp_path: Path) -> None:
    from sepsis_mcp.appendix_score_diagnostics import (
        GossisScoreDiagnosticsConfig,
        MimicScoreDiagnosticsConfig,
        run_score_diagnostics,
    )

    gossis_root = write_structured_synthetic_gossis_dataset(tmp_path)
    mimic_csv = _write_synthetic_mimic4_csv(tmp_path)
    output_dir = tmp_path / "score-diagnostics"

    paths = run_score_diagnostics(
        output_dir=output_dir,
        gossis_config=GossisScoreDiagnosticsConfig(
            data_root=gossis_root,
            min_hospital_admissions=100,
            random_state_grid=(0,),
            model_type_grid=("xgboost",),
        ),
        mimic_config=MimicScoreDiagnosticsConfig(
            csv_path=mimic_csv,
            min_unit_size=10,
            rotations=1,
            max_assignments=1,
            min_selection_group_rows=2,
            model_type_grid=("xgboost",),
        ),
    )

    assert Path(paths["score_records"]).exists()
    assert Path(paths["score_summary"]).exists()
    assert Path(paths["score_figure"]).exists()

    score_records = pd.read_csv(paths["score_records"])
    score_summary = pd.read_csv(paths["score_summary"])

    assert {"dataset", "model_type", "seed", "split_kind", "group", "group_label", "score"} <= set(score_records.columns)
    assert {"dataset", "model_type", "group_label", "pooled_ks", "calibration_vs_test_ks"} <= set(score_summary.columns)
    assert set(score_records["dataset"]) == {"gossis", "mimic4"}
    assert set(score_records["split_kind"]) == {"calibration", "test"}