src/workflows/survival.py

"""
Survival analysis workflow helpers (placeholder module).

Provides scaffolding for survival analysis using activity scores or subtype
labels derived from the decoupleRpy pipeline.  All functions here are
placeholders — implementation requires:

  1. Validated survival metadata (event + time columns) in dataset manifests.
  2. A survival analysis dependency added to requirements.txt
     (lifelines or scikit-survival).
  3. Confirmation that the target datasets include survival endpoints.

Functions
---------
check_survival_data_available  Inspect a manifest for survival metadata.
run_kaplan_meier               Placeholder: KM curves stratified by group.
run_cox_regression             Placeholder: Cox PH model with activity covariates.

Dataset-specific note
---------------------
GSE71729 (Moffitt) does not include survival endpoints in the GEO series
matrix deposit.  Survival data for this cohort would need to be obtained
from the supplementary tables of Moffitt et al. 2015 and added to the
manifest's obs.other_columns before these functions can be used.
"""

from typing import Any


def check_survival_data_available(manifest: dict) -> dict[str, Any]:
    """
    Inspect a manifest for survival metadata (event and time columns).

    Looks for column names containing "event", "status", "time", or "survival"
    in the manifest's obs.other_columns list.

    Returns
    -------
    dict with keys:
        available     (bool)      — True if both event and time columns found.
        event_column  (str|None)  — Detected event/status column name.
        time_column   (str|None)  — Detected time/survival column name.
        message       (str)       — Human-readable summary.
    """
    # Support both DatasetManifest dataclass and raw dict
    if hasattr(manifest, "survival_columns"):
        sc = manifest.survival_columns or {}
        if sc.get("event_column") and sc.get("time_column"):
            return {
                "available": True,
                "event_column": sc["event_column"],
                "time_column": sc["time_column"],
                "message": "survival columns found in manifest.survival_columns",
            }
        return {
            "available": False,
            "event_column": None,
            "time_column": None,
            "message": "survival_columns not set or null in manifest",
        }

    obs = manifest.get("obs", {})
    other_cols: list[str] = obs.get("other_columns", [])

    event_col = next(
        (c for c in other_cols if any(k in c.lower() for k in ("event", "status"))),
        None,
    )
    time_col = next(
        (c for c in other_cols if any(k in c.lower() for k in ("time", "survival", "os_"))),
        None,
    )

    available = event_col is not None and time_col is not None
    message = (
        f"survival columns found: event='{event_col}', time='{time_col}'"
        if available
        else "survival columns not declared in manifest obs.other_columns"
    )

    return {
        "available": available,
        "event_column": event_col,
        "time_column": time_col,
        "message": message,
    }


def run_kaplan_meier(
    adata_path: str,
    manifest: dict,
    group_column: str | None = None,
    out_prefix: str | None = None,
) -> dict:
    """
    Generate Kaplan-Meier survival curves stratified by a group column.

    NOT YET IMPLEMENTED — deferred until:
    - Survival metadata is confirmed in the target dataset's manifest.
    - lifelines is added to requirements.txt.

    Parameters
    ----------
    adata_path:    Path to h5ad with sample-level obs including survival columns.
    manifest:      Parsed dataset manifest dict.
    group_column:  Column for stratification (defaults to obs.condition_column).
    out_prefix:    Output file prefix.
    """
    survival_info = check_survival_data_available(manifest)
    if not survival_info["available"]:
        raise ValueError(
            f"Cannot run Kaplan-Meier: {survival_info['message']}. "
            "Add event and time columns to the manifest's obs.other_columns."
        )
    raise NotImplementedError("run_kaplan_meier is a placeholder. Requires lifelines.")


def run_cox_regression(
    activity_scores_path: str,
    manifest: dict,
    covariates: list[str] | None = None,
    out_prefix: str | None = None,
) -> dict:
    """
    Fit a Cox proportional hazards model with activity scores as covariates.

    NOT YET IMPLEMENTED — deferred until survival metadata and a survival
    analysis library are confirmed available.

    Parameters
    ----------
    activity_scores_path:  Path to activity scores CSV (samples × features).
    manifest:              Parsed dataset manifest dict.
    covariates:            Feature names to include as covariates.
                           None = use all features in the scores file.
    out_prefix:            Output file prefix.
    """
    # Params are part of the planned signature; intentionally unused until the
    # Cox model is implemented (see NotImplementedError below).
    _ = (activity_scores_path, covariates, out_prefix)
    survival_info = check_survival_data_available(manifest)
    if not survival_info["available"]:
        raise ValueError(f"Cannot run Cox regression: {survival_info['message']}.")
    raise NotImplementedError("run_cox_regression is a placeholder. Requires lifelines.")