src/survival_utils.py

import numpy as np
import pandas as pd
from lifelines import CoxPHFitter

DEFAULT_YEARS = [3, 5, 10]


def build_os_time_days(
    df: pd.DataFrame,
    months_col: str = "Months_from_transplant_to_last_followup",
    death_date_col: str = "Date_of_death",
    hsct_date_col: str = "HSCT_date",
    event_col: str = "Event",
) -> pd.Series:
    """
    Build OS_time_days with priority:
    1) Months_from_transplant_to_last_followup (converted to days)
    2) If missing, and Event == 1, use Date_of_death - HSCT_date

    Returns
    -------
    pd.Series
        Survival duration in days.
    """
    out = pd.Series(np.nan, index=df.index, dtype="float64")

    # First preference: follow-up months already present
    if months_col in df.columns:
        months = pd.to_numeric(df[months_col], errors="coerce")
        valid_months = months.notna() & (months > 0)
        out.loc[valid_months] = months.loc[valid_months] * 30.44

    # Fallback: calculate from death date - HSCT date for death events
    required_cols = [death_date_col, hsct_date_col, event_col]
    if all(c in df.columns for c in required_cols):
        death_dt = pd.to_datetime(df[death_date_col], errors="coerce")
        hsct_dt = pd.to_datetime(df[hsct_date_col], errors="coerce")
        event = pd.to_numeric(df[event_col], errors="coerce")

        mask = out.isna() & death_dt.notna() & hsct_dt.notna() & (event == 1)
        out.loc[mask] = (death_dt.loc[mask] - hsct_dt.loc[mask]).dt.days.astype(float)

    # Clean invalid durations
    out.loc[out <= 0] = np.nan
    return out


def prepare_cox_df(
    df_surv: pd.DataFrame,
    preds_surv: np.ndarray,
    covariates: list[str],
    duration_col: str = "OS_time_days",
    event_col: str = "Event_clean",
    risk_col_name: str = "Predicted_GVHD_Risk",
) -> tuple[pd.DataFrame, list[str], list[str]]:
    """
    Prepare Cox dataframe for OS modeling.

    Parameters
    ----------
    df_surv : pd.DataFrame
        Survival-ready dataframe, usually preprocessed dataframe.
    preds_surv : np.ndarray
        Predicted risk scores aligned with df_surv rows.
    covariates : list[str]
        Covariates to include in Cox model.
    duration_col : str
        Name of duration column to use/create.
    event_col : str
        Name of cleaned event column to use/create.
    risk_col_name : str
        Explicit risk predictor name, e.g.:
        - Predicted_acute_gvhd_risk
        - Predicted_chronic_gvhd_risk

    Returns
    -------
    df_cox : pd.DataFrame
        Numeric-only Cox dataset including duration and event columns.
    design_cols : list[str]
        Final model covariates after one-hot encoding.
    cat_cols : list[str]
        Original categorical covariates that were one-hot encoded.
    """
    df = df_surv.copy()

    # Build OS_time_days if not already present or if missing
    if duration_col not in df.columns:
        df[duration_col] = build_os_time_days(df)
    else:
        existing_duration = pd.to_numeric(df[duration_col], errors="coerce")
        rebuilt_duration = build_os_time_days(df)
        df[duration_col] = existing_duration
        fill_mask = df[duration_col].isna()
        df.loc[fill_mask, duration_col] = rebuilt_duration.loc[fill_mask]

    # Build Event_clean if not already present
    if event_col not in df.columns:
        if "Event" not in df.columns:
            raise ValueError(
                f"Neither '{event_col}' nor raw 'Event' column found in dataframe."
            )
        df[event_col] = pd.to_numeric(df["Event"], errors="coerce")
    else:
        df[event_col] = pd.to_numeric(df[event_col], errors="coerce")

    # Add predicted GVHD risk as Cox covariate
    preds_surv = np.asarray(preds_surv).astype(float).ravel()
    if len(preds_surv) != len(df):
        raise ValueError(
            f"Length mismatch: preds_surv has {len(preds_surv)} rows but df_surv has {len(df)} rows."
        )
    df[risk_col_name] = preds_surv

    # Keep only required columns
    keep_covariates = [c for c in covariates if c in df.columns]
    cols = [duration_col, event_col, risk_col_name] + keep_covariates
    df = df[cols].copy()

    # Coerce survival columns
    df[duration_col] = pd.to_numeric(df[duration_col], errors="coerce")
    df[event_col] = pd.to_numeric(df[event_col], errors="coerce")

    # Drop invalid rows
    df = df.dropna(subset=[duration_col, event_col])
    df = df[df[duration_col] > 0]
    df[event_col] = df[event_col].astype(int).clip(0, 1)

    # Identify categorical covariates
    from pandas.api.types import is_object_dtype, is_string_dtype, is_categorical_dtype

    cat_cols = [
        c for c in df.columns
        if c not in [duration_col, event_col]
        and (is_object_dtype(df[c]) or is_string_dtype(df[c]) or is_categorical_dtype(df[c]))
    ]

    # One-hot encode categoricals
    if cat_cols:
        df = pd.get_dummies(df, columns=cat_cols, drop_first=True)

    # Ensure numeric design matrix
    for c in df.columns:
        if c in [duration_col, event_col]:
            continue
        df[c] = pd.to_numeric(df[c], errors="coerce").fillna(0)

    design_cols = [c for c in df.columns if c not in [duration_col, event_col]]
    return df, design_cols, cat_cols


def fit_cox(
    df_cox: pd.DataFrame,
    duration_col: str = "OS_time_days",
    event_col: str = "Event_clean",
) -> CoxPHFitter:
    cph = CoxPHFitter(penalizer=0.1)
    cph.fit(df_cox, duration_col=duration_col, event_col=event_col)
    return cph


def make_patient_design_row(
    patient_row: pd.DataFrame,
    design_cols: list[str],
    cat_cols_original: list[str],
):
    x = patient_row.copy()

    if cat_cols_original:
        available_cat_cols = [c for c in cat_cols_original if c in x.columns]
        if available_cat_cols:
            x = pd.get_dummies(x, columns=available_cat_cols, drop_first=True)

    unseen = sorted(set(x.columns) - set(design_cols))
    x = x.reindex(columns=design_cols, fill_value=0)

    for c in x.columns:
        x[c] = pd.to_numeric(x[c], errors="coerce").fillna(0)

    return x, unseen


def predict_patient_survival(cph, patient_design, years=None):
    if years is None:
        years = DEFAULT_YEARS

    surv_fn = cph.predict_survival_function(patient_design)

    # survival function is a DataFrame: index=time, column=patient
    s_series = surv_fn.iloc[:, 0]
    times = s_series.index.values

    landmarks = {}
    for y in years:
        t = 365 * int(y)

        le_mask = times <= t
        if le_mask.any():
            s = float(s_series.loc[times[le_mask].max()])
        else:
            s = float(s_series.iloc[0])

        landmarks[int(y)] = s

    return surv_fn, landmarks