src/workflows/activity_stats.py

"""
Statistical analysis of TF and pathway activity scores.

Provides dataset-agnostic functions for extracting contrast definitions
from manifests, ranking activity scores, and comparing activities between
metadata-defined sample groups.

Functions (implemented)
-----------------------
get_contrast_groups        Extract test/control group labels from a manifest.
rank_by_activity           Rank TFs or pathways by mean absolute activity score.
compare_activity_by_group  Welch's t-test + BH FDR across all activities between
                           two metadata-defined sample groups.  Returns a table
                           with activity, mean_test, mean_control, effect_size
                           (Cohen's d), statistic, pvalue, padj, n_test, n_control.

"""

from __future__ import annotations

from typing import Any

import numpy as np
import pandas as pd

from src.workflows.metadata_validation import subset_and_require_group


def get_contrast_groups(manifest: dict | object) -> dict[str, str]:
    """
    Extract the contrast definition from a manifest.

    Accepts both a DatasetManifest dataclass instance and a raw dict so this
    function works regardless of how the manifest was loaded.

    Returns a dict with keys:
        design_factor (str)  — column in obs used as the grouping variable.
        test_group    (str)  — label of the test / foreground condition.
        control_group (str)  — label of the reference / background condition.
        method        (str)  — DE method: "deseq2", "ttest", or "limma".

    Returns empty strings for any missing fields so callers can check
    truthiness rather than catching KeyError.
    """
    if hasattr(manifest, "default_contrasts"):
        contrasts = manifest.default_contrasts or []
        contrast = contrasts[0] if contrasts else {}
    else:
        contrast = manifest.get("contrast", {})

    return {
        "design_factor": contrast.get("design_factor", "condition"),
        "test_group": contrast.get("test_group", ""),
        "control_group": contrast.get("control_group", ""),
        "method": contrast.get("method", "ttest"),
    }


def rank_by_activity(
    activities: dict[str, float],
    top_n: int = 25,
    by_abs: bool = True,
) -> list[tuple[str, float]]:
    """
    Rank features (TFs, pathways) by activity score.

    Parameters
    ----------
    activities:  Mapping of feature_name → activity_score (float).
    top_n:       Number of top features to return.
    by_abs:      If True, rank by |score| so both activated and repressed
                 features surface.  If False, rank by raw score (top = most
                 activated).

    Returns
    -------
    Sorted list of (feature_name, score) tuples, descending by the ranking key.
    """
    key = (lambda kv: abs(kv[1])) if by_abs else (lambda kv: kv[1])
    ranked = sorted(activities.items(), key=key, reverse=True)
    return ranked[:top_n]


def compare_activity_by_group(
    activity_df: pd.DataFrame,
    metadata_df: pd.DataFrame,
    group_column: str,
    test_group: str,
    control_group: str,
    subset_query: str | None = None,
    method: str = "welch_ttest",
) -> dict[str, Any]:
    """
    Compare activity scores between two metadata-defined sample groups.

    Use this after dataset_score_bulk_samples when the user asks whether
    TF/pathway/hallmark activity differs between two sample groups (e.g.
    Classical vs Basal PDAC, treated vs untreated, tumour vs normal).

    Runs Welch's t-test independently for each activity (column) and applies
    Benjamini-Hochberg FDR correction across all tests.  Effect size is
    Cohen's d: (mean_test − mean_control) / pooled_std — positive values
    indicate higher activity in the test group.

    This function operates on pre-computed activity scores (output of
    score_bulk_samples_with_decoupler), not on raw expression matrices.

    Parameters
    ----------
    activity_df:
        Samples × activities DataFrame (e.g. output of
        score_bulk_samples_with_decoupler).  Index = sample IDs.
    metadata_df:
        Sample metadata DataFrame.  Index must align with activity_df.index.
    group_column:
        Column in metadata_df containing the group labels.
    test_group:
        Label of the foreground / test condition.
    control_group:
        Label of the reference / background condition.
    subset_query:
        Optional pandas query string applied to metadata_df before grouping,
        e.g. ``"tissue == 'tumor'"``.
    method:
        Statistical method.  Currently only "welch_ttest" is supported.

    Returns
    -------
    dict with keys:
        dataframe       (pd.DataFrame) — activities × stats, sorted by padj.
                        Columns: mean_test, mean_control, effect_size,
                                 statistic, pvalue, padj, n_test, n_control.
                        Index name: "activity".
        n_activities    (int)
        n_test          (int)
        n_control       (int)
        group_column    (str)
        test_group      (str)
        control_group   (str)
        subset_query    (str | None)
        method          (str)
        significant_05  (int) — activities with padj < 0.05
        significant_01  (int) — activities with padj < 0.01
        warnings        (list[str])

    Raises
    ------
    ValueError  if group_column is missing, groups are not found, either group
                has fewer than 2 samples, or subset_query fails.
    """
    import warnings as _w

    from scipy import stats
    from statsmodels.stats.multitest import multipletests

    _SUPPORTED = frozenset({"welch_ttest"})
    if method not in _SUPPORTED:
        raise ValueError(f"method must be one of {sorted(_SUPPORTED)}, got '{method}'")

    run_warnings: list[str] = []

    # Subset (optional) + require the group column (shared helper).
    working_meta = subset_and_require_group(metadata_df, subset_query, group_column)

    # ── Validate group labels ─────────────────────────────────────────────
    available = working_meta[group_column].unique().tolist()
    available_str = sorted(str(g) for g in available)

    if test_group == control_group:
        raise ValueError("test_group and control_group must be different.")
    if test_group not in available:
        raise ValueError(
            f"test_group '{test_group}' not found in column '{group_column}'. "
            f"Available: {available_str}"
        )
    if control_group not in available:
        raise ValueError(
            f"control_group '{control_group}' not found in column '{group_column}'. "
            f"Available: {available_str}"
        )

    # ── Align activity_df to (subsetted) metadata ─────────────────────────
    common_idx = activity_df.index.intersection(working_meta.index)
    if len(common_idx) == 0:
        raise ValueError(
            "No common samples between activity_df and metadata_df. "
            "Check that indices are aligned (both should use sample IDs)."
        )
    if len(common_idx) < len(activity_df):
        run_warnings.append(
            f"{len(activity_df) - len(common_idx)} activity sample(s) had no "
            "metadata match and were excluded."
        )

    aligned_act = activity_df.loc[common_idx]
    aligned_meta = working_meta.loc[common_idx]

    # ── Build group masks ─────────────────────────────────────────────────
    test_mask = aligned_meta[group_column] == test_group
    ctrl_mask = aligned_meta[group_column] == control_group
    n_test = int(test_mask.sum())
    n_ctrl = int(ctrl_mask.sum())

    if n_test < 2:
        raise ValueError(
            f"test_group '{test_group}' has {n_test} sample(s) after alignment — "
            "need at least 2 for Welch's t-test."
        )
    if n_ctrl < 2:
        raise ValueError(
            f"control_group '{control_group}' has {n_ctrl} sample(s) after alignment — "
            "need at least 2 for Welch's t-test."
        )
    if n_test < 5 or n_ctrl < 5:
        run_warnings.append(
            f"Small group sizes (test={n_test}, control={n_ctrl}). "
            "Statistical power is limited; interpret results with caution."
        )

    n_excluded = len(aligned_act) - n_test - n_ctrl
    if n_excluded > 0:
        other = [str(g) for g in available if g not in (test_group, control_group)]
        run_warnings.append(
            f"{n_excluded} sample(s) in other group(s) excluded from comparison "
            f"({other[:5]}{'...' if len(other) > 5 else ''})."
        )

    # ── Vectorised Welch's t-test ─────────────────────────────────────────
    X_test = aligned_act[test_mask].values  # n_test × n_activities
    X_ctrl = aligned_act[ctrl_mask].values  # n_ctrl × n_activities

    with _w.catch_warnings():
        _w.simplefilter("ignore", RuntimeWarning)
        t_stats, p_vals = stats.ttest_ind(X_test, X_ctrl, axis=0, equal_var=False)

    # Replace NaN / Inf from zero-variance activities before BH correction
    t_stats = np.where(np.isfinite(t_stats), t_stats, 0.0)
    p_vals = np.where(np.isfinite(p_vals), p_vals, 1.0)

    _, padj, _, _ = multipletests(p_vals, method="fdr_bh")

    mean_test = X_test.mean(axis=0)
    mean_ctrl = X_ctrl.mean(axis=0)

    # ── Cohen's d effect size ─────────────────────────────────────────────
    # pooled_std = sqrt(((n1-1)*s1² + (n2-1)*s2²) / (n1+n2-2))
    # Positive d → test_group has higher activity.
    std_test = X_test.std(axis=0, ddof=1)
    std_ctrl = X_ctrl.std(axis=0, ddof=1)
    pooled_var = ((n_test - 1) * std_test**2 + (n_ctrl - 1) * std_ctrl**2) / (n_test + n_ctrl - 2)
    with np.errstate(divide="ignore", invalid="ignore"):
        cohens_d = np.where(
            pooled_var > 0,
            (mean_test - mean_ctrl) / np.sqrt(pooled_var),
            0.0,
        )

    activities = aligned_act.columns.tolist()
    result_df = pd.DataFrame(
        {
            "mean_test": mean_test,
            "mean_control": mean_ctrl,
            "effect_size": cohens_d,
            "statistic": t_stats,
            "pvalue": p_vals,
            "padj": padj,
            "n_test": n_test,
            "n_control": n_ctrl,
        },
        index=activities,
    )
    result_df.index.name = "activity"
    result_df = result_df.sort_values("padj")

    run_warnings.append(
        f"effect_size is Cohen's d: (mean_test − mean_control) / pooled_std. "
        f"Positive values indicate higher activity in '{test_group}'. "
        "Activities with zero variance in either group are assigned d=0."
    )

    return {
        "dataframe": result_df,
        "n_activities": len(activities),
        "n_test": n_test,
        "n_control": n_ctrl,
        "group_column": group_column,
        "test_group": test_group,
        "control_group": control_group,
        "subset_query": subset_query,
        "method": method,
        "significant_05": int((padj < 0.05).sum()),
        "significant_01": int((padj < 0.01).sum()),
        "warnings": run_warnings,
    }