src/sepsis_mcp/missingness_groups.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import numpy as np
import pandas as pd
from sklearn.cluster import KMeans
from sklearn.decomposition import TruncatedSVD
from sklearn.metrics import silhouette_score


@dataclass(frozen=True)
class MissingnessGroupingResult:
    selection_groups: pd.DataFrame
    calibration_groups: pd.DataFrame
    test_groups: pd.DataFrame
    group_labels: dict[int, str]
    metadata: dict[str, Any]


def build_missingness_groups(
    *,
    selection_frame: pd.DataFrame,
    calibration_frame: pd.DataFrame,
    test_frame: pd.DataFrame,
    strategy: str,
    candidate_missing_variables: list[str] | None = None,
    mask_cluster_k_grid: list[int] | None = None,
    min_group_fraction: float = 0.10,
    min_selection_group_rows: int = 1,
    random_state: int = 0,
    **_: Any,
) -> MissingnessGroupingResult:
    if strategy == "coverage_gap_variable":
        return _build_coverage_gap_variable_groups(
            selection_frame=selection_frame,
            calibration_frame=calibration_frame,
            test_frame=test_frame,
            candidate_missing_variables=candidate_missing_variables,
            min_group_fraction=min_group_fraction,
            min_selection_group_rows=min_selection_group_rows,
        )
    if strategy == "mask_cluster":
        return _build_mask_cluster_groups(
            selection_frame=selection_frame,
            calibration_frame=calibration_frame,
            test_frame=test_frame,
            mask_cluster_k_grid=mask_cluster_k_grid or [2, 3, 4, 5],
            min_group_fraction=min_group_fraction,
            min_selection_group_rows=min_selection_group_rows,
            random_state=random_state,
        )
    raise ValueError(f"unsupported grouping strategy: {strategy}")


def _missing_rate_columns(frame: pd.DataFrame) -> list[str]:
    return [
        column
        for column in frame.columns
        if column.endswith("_missing_rate") and column != "global_missing_rate"
    ]


def _candidate_missing_rate_columns(
    frame: pd.DataFrame,
    candidate_missing_variables: list[str] | None,
) -> list[str]:
    available = set(_missing_rate_columns(frame))
    if candidate_missing_variables is None:
        return sorted(available)
    columns = []
    for variable in candidate_missing_variables:
        column = f"{variable}_missing_rate"
        if column in available:
            columns.append(column)
    return columns


def _never_observed_indicator(frame: pd.DataFrame, column: str) -> np.ndarray:
    values = frame[column].to_numpy(dtype=float)
    return values > 0.0


def _frame_groups_from_ids(group_ids: np.ndarray, group_labels: dict[int, str]) -> pd.DataFrame:
    return pd.DataFrame(
        {
            "group": group_ids.astype(int),
            "group_label": [group_labels[int(group_id)] for group_id in group_ids],
        }
    )


def _build_coverage_gap_variable_groups(
    *,
    selection_frame: pd.DataFrame,
    calibration_frame: pd.DataFrame,
    test_frame: pd.DataFrame,
    candidate_missing_variables: list[str] | None,
    min_group_fraction: float,
    min_selection_group_rows: int,
) -> MissingnessGroupingResult:
    candidate_columns = _candidate_missing_rate_columns(selection_frame, candidate_missing_variables)
    diagnostics: dict[str, dict[str, float | int]] = {}
    candidate_order = {column: index for index, column in enumerate(candidate_columns)}

    def collect_diagnostics(*, relaxed: bool) -> dict[str, dict[str, float | int]]:
        local_diagnostics: dict[str, dict[str, float | int]] = {}
        for column in candidate_columns:
            selection_missing = _never_observed_indicator(selection_frame, column)
            calibration_missing = _never_observed_indicator(calibration_frame, column)

            selection_fraction = float(selection_missing.mean())
            if not relaxed and (
                selection_fraction < min_group_fraction
                or selection_fraction > (1.0 - min_group_fraction)
            ):
                continue

            selection_missing_rows = int(selection_missing.sum())
            selection_observed_rows = int((~selection_missing).sum())
            calibration_missing_rows = int(calibration_missing.sum())
            calibration_observed_rows = int((~calibration_missing).sum())
            if not relaxed and min(selection_missing_rows, selection_observed_rows) < min_selection_group_rows:
                continue
            calibration_fraction = float(calibration_missing.mean())
            selection_gap = abs(selection_fraction - calibration_fraction)
            imbalance = abs(0.5 - selection_fraction)
            local_diagnostics[column] = {
                "selection_missing_fraction": selection_fraction,
                "calibration_missing_fraction": calibration_fraction,
                "selection_missing_rows": selection_missing_rows,
                "selection_observed_rows": selection_observed_rows,
                "calibration_missing_rows": calibration_missing_rows,
                "calibration_observed_rows": calibration_observed_rows,
                "selection_gap": selection_gap,
                "imbalance": imbalance,
                "minority_support": min(
                    selection_missing_rows,
                    selection_observed_rows,
                ),
                "relaxed_selection": relaxed,
            }
        return local_diagnostics

    diagnostics = collect_diagnostics(relaxed=False)
    if not diagnostics and candidate_missing_variables is None:
        diagnostics = collect_diagnostics(relaxed=True)
    if not diagnostics:
        raise ValueError("no candidates satisfied minimum support and missing-fraction requirements")

    for column in candidate_columns:
        selection_missing = _never_observed_indicator(selection_frame, column)
        calibration_missing = _never_observed_indicator(calibration_frame, column)

        selection_fraction = float(selection_missing.mean())
        if (
            selection_fraction < min_group_fraction
            or selection_fraction > (1.0 - min_group_fraction)
        ):
            continue

        selection_missing_rows = int(selection_missing.sum())
        selection_observed_rows = int((~selection_missing).sum())
        calibration_missing_rows = int(calibration_missing.sum())
        calibration_observed_rows = int((~calibration_missing).sum())
        if min(selection_missing_rows, selection_observed_rows) < min_selection_group_rows:
            continue

        calibration_fraction = float(calibration_missing.mean())
        selection_gap = abs(selection_fraction - calibration_fraction)
        imbalance = abs(0.5 - selection_fraction)
        if column not in diagnostics:
            continue
        diagnostics[column].update(
            {
                "selection_missing_fraction": selection_fraction,
                "calibration_missing_fraction": calibration_fraction,
                "selection_missing_rows": selection_missing_rows,
                "selection_observed_rows": selection_observed_rows,
                "calibration_missing_rows": calibration_missing_rows,
                "calibration_observed_rows": calibration_observed_rows,
                "selection_gap": selection_gap,
                "imbalance": imbalance,
                "minority_support": min(
                    selection_missing_rows,
                    selection_observed_rows,
                ),
            }
        )

    selected_column = sorted(
        diagnostics,
        key=lambda column: (
            -float(diagnostics[column]["selection_gap"]),
            float(diagnostics[column]["imbalance"]),
            -int(diagnostics[column]["minority_support"]),
            candidate_order.get(column, len(candidate_order)),
        ),
    )[0]

    variable = selected_column[: -len("_missing_rate")]
    group_labels = {
        0: f"{variable.lower()}_ever_observed",
        1: f"{variable.lower()}_never_observed",
    }

    selection_ids = _never_observed_indicator(selection_frame, selected_column).astype(int)
    calibration_ids = _never_observed_indicator(calibration_frame, selected_column).astype(int)
    test_ids = _never_observed_indicator(test_frame, selected_column).astype(int)
    return MissingnessGroupingResult(
        selection_groups=_frame_groups_from_ids(selection_ids, group_labels),
        calibration_groups=_frame_groups_from_ids(calibration_ids, group_labels),
        test_groups=_frame_groups_from_ids(test_ids, group_labels),
        group_labels=group_labels,
        metadata={
            "strategy": "coverage_gap_variable",
            "selected_variable": variable,
            "group_source": "missing_rate_column",
            "group_count": len(group_labels),
            "selection_variable_diagnostics": {
                column[: -len("_missing_rate")]: values for column, values in diagnostics.items()
            },
        },
    )


def _mask_feature_matrix(frame: pd.DataFrame, columns: list[str]) -> np.ndarray:
    return np.column_stack([_never_observed_indicator(frame, column).astype(float) for column in columns])


def _filter_mask_columns(
    selection_frame: pd.DataFrame,
    columns: list[str],
    *,
    min_group_fraction: float,
) -> list[str]:
    filtered = []
    for column in columns:
        fraction = float(_never_observed_indicator(selection_frame, column).mean())
        if min_group_fraction <= fraction <= (1.0 - min_group_fraction):
            filtered.append(column)
    return filtered


def _build_mask_cluster_groups(
    *,
    selection_frame: pd.DataFrame,
    calibration_frame: pd.DataFrame,
    test_frame: pd.DataFrame,
    mask_cluster_k_grid: list[int],
    min_group_fraction: float,
    min_selection_group_rows: int,
    random_state: int,
) -> MissingnessGroupingResult:
    columns = _filter_mask_columns(
        selection_frame,
        _missing_rate_columns(selection_frame),
        min_group_fraction=min_group_fraction,
    )
    if len(columns) < 2:
        raise ValueError("mask clustering requires at least two non-constant missingness columns")

    selection_mask = _mask_feature_matrix(selection_frame, columns)
    calibration_mask = _mask_feature_matrix(calibration_frame, columns)
    test_mask = _mask_feature_matrix(test_frame, columns)
    n_components = min(10, selection_mask.shape[0], selection_mask.shape[1])
    if n_components >= 1 and n_components < selection_mask.shape[1]:
        reducer = TruncatedSVD(n_components=n_components, random_state=random_state)
        selection_features = reducer.fit_transform(selection_mask)
        calibration_features = reducer.transform(calibration_mask)
        test_features = reducer.transform(test_mask)
    else:
        selection_features = selection_mask
        calibration_features = calibration_mask
        test_features = test_mask

    diagnostics: dict[int, dict[str, float | int]] = {}
    best_k = None
    best_score = None
    best_model = None
    for k in mask_cluster_k_grid:
        if k <= 1 or k > len(selection_features):
            diagnostics[k] = {"silhouette": float("nan"), "min_cluster_size": 0, "selection_gap": float("nan")}
            continue
        model = KMeans(n_clusters=k, random_state=random_state, n_init=10)
        selection_ids = model.fit_predict(selection_features)
        counts = np.bincount(selection_ids, minlength=k)
        min_cluster_size = int(counts.min()) if counts.size else 0
        if min_cluster_size < min_selection_group_rows:
            diagnostics[k] = {"silhouette": float("nan"), "min_cluster_size": min_cluster_size, "selection_gap": float("nan")}
            continue
        silhouette = (
            float(silhouette_score(selection_features, selection_ids))
            if len(np.unique(selection_ids)) > 1
            else float("nan")
        )
        cluster_means = []
        for cluster_id in range(k):
            cluster_mask = selection_ids == cluster_id
            cluster_means.append(float(selection_frame.loc[cluster_mask, "global_missing_rate"].mean()))
        selection_gap = float(max(cluster_means) - min(cluster_means)) if cluster_means else 0.0
        diagnostics[k] = {
            "silhouette": silhouette,
            "min_cluster_size": min_cluster_size,
            "selection_gap": selection_gap,
        }
        score = (-k,)
        if best_score is None or score > best_score:
            best_score = score
            best_k = k
            best_model = model

    if best_model is None or best_k is None:
        raise ValueError("mask clustering could not find a feasible k")

    selection_ids = best_model.predict(selection_features)
    calibration_ids = best_model.predict(calibration_features)
    test_ids = best_model.predict(test_features)
    group_labels = {group_id: f"cluster_{group_id}" for group_id in range(best_k)}
    return MissingnessGroupingResult(
        selection_groups=_frame_groups_from_ids(selection_ids, group_labels),
        calibration_groups=_frame_groups_from_ids(calibration_ids, group_labels),
        test_groups=_frame_groups_from_ids(test_ids, group_labels),
        group_labels=group_labels,
        metadata={
            "strategy": "mask_cluster",
            "selected_k": best_k,
            "group_source": "missingness_mask_cluster",
            "group_count": len(group_labels),
            "mask_cluster_diagnostics": diagnostics,
        },
    )


__all__ = ["MissingnessGroupingResult", "build_missingness_groups"]