src/workflows/activity_scoring.py

"""
Transcription factor and pathway activity scoring workflow helpers.

Two families of functions live here:

1. Manifest-aware parameter helpers
   ----------------------------------
   get_scoring_params   Extract organism, contrast, and applicable workflows
                        from a dataset manifest.

2. Generic sample-level scoring
   ------------------------------
   score_bulk_samples_with_decoupler
                        Run decoupleR activity estimation (ULM/MLM/zscore) on a
                        pre-normalised samples × genes expression DataFrame.
                        Returns a samples × activities DataFrame and diagnostic
                        metadata.  Does not hard-code any dataset.
"""

from __future__ import annotations

from typing import Any

import pandas as pd

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

SUPPORTED_RESOURCES: frozenset[str] = frozenset({"progeny", "collectri", "hallmark"})

SUPPORTED_METHODS: frozenset[str] = frozenset({"ulm", "mlm", "zscore"})

_RESOURCE_DESCRIPTIONS: dict[str, str] = {
    "progeny":   "PROGENy — 14 cancer signalling pathway signatures",
    "collectri": "CollecTRI — TF–target regulons (~1 000 TFs)",
    "hallmark":  "MSigDB Hallmark — 50 curated gene sets",
}


# ---------------------------------------------------------------------------
# Manifest helpers
# ---------------------------------------------------------------------------

def get_scoring_params(manifest: "dict | object") -> dict[str, Any]:
    """
    Extract activity scoring parameters from a dataset manifest.

    Accepts both a DatasetManifest dataclass instance and a raw dict.

    Returns
    -------
    dict with keys:
        organism  (str)  — "human" or "mouse".
        contrast  (dict) — design_factor, test_group, control_group, method.
        workflows (list) — which workflow modules apply to this dataset.
    """
    if hasattr(manifest, "organism"):
        return {
            "organism": manifest.organism,
            "contrast": (manifest.default_contrasts or [{}])[0]
                        if hasattr(manifest, "default_contrasts") else {},
            "workflows": getattr(manifest, "valid_workflows", []),
        }
    return {
        "organism": manifest.get("organism", "human"),
        "contrast": manifest.get("contrast", {}),
        "workflows": manifest.get("workflows", []),
    }


# ---------------------------------------------------------------------------
# Internal helper
# ---------------------------------------------------------------------------

def _find_network_target_col(net: "pd.DataFrame") -> "str | None":
    """Return the target/gene column name in a decoupler network, or None."""
    for col in ("target", "gene", "Gene", "to"):
        if col in net.columns:
            return col
    return None


# ---------------------------------------------------------------------------
# Generic sample-level scoring
# ---------------------------------------------------------------------------

def score_bulk_samples_with_decoupler(
    expression_df: "pd.DataFrame",
    resource: str = "progeny",
    organism: str = "human",
    method: str = "ulm",
    min_n: int = 5,
    _network: "pd.DataFrame | None" = None,
) -> dict[str, Any]:
    """
    Score bulk expression samples using decoupleR activity estimation.

    Computes per-sample activity scores for pathways (PROGENy), TF regulons
    (CollecTRI), or hallmark gene sets by fitting a linear model (ULM/MLM)
    or weighted sum (WSUM) across the genes in each regulon/pathway.

    This function is generic — it does not assume any specific dataset.
    Pass any pre-normalised samples × genes expression matrix.

    Input requirements
    ------------------
    expression_df:
        Samples as rows (index = sample IDs), genes as columns (HGNC symbols).
        Values must be pre-normalised (log-CPM, log-intensity, log2-TPM, etc.).
        Do NOT pass raw integer counts — activity estimation assumes continuous,
        roughly symmetric expression values.

    Output
    ------
    activity_df:
        Same row index as expression_df.  Columns are activity names
        (pathway names, TF names, or hallmark names).

    Parameters
    ----------
    resource:
        "progeny"   — PROGENy pathway signatures (14 pathways).
        "collectri" — CollecTRI TF regulons (~1 000 TFs, human/mouse).
        "hallmark"  — MSigDB Hallmark gene sets (50 sets).
    organism:
        "human" or "mouse".
    method:
        "ulm"  — Univariate Linear Model (recommended; fastest).
        "mlm"  — Multivariate Linear Model (accounts for co-linearity).
        "zscore" — Z-score (simpler, no distributional assumptions).
    min_n:
        Minimum expected samples. Emits a warning when fewer are present.
    _network:
        For testing only.  Supply a pre-loaded network DataFrame to bypass
        the decoupler network download (dc.op.*).  Not exposed in MCP tools.

    Returns
    -------
    dict with keys:
        activity_df     (pd.DataFrame) — samples × activities.
        pvalue_df       (pd.DataFrame) — samples × activities (p-values).
        n_samples       (int)
        n_activities    (int)   — number of scored activities (sources).
        n_network_genes (int)   — total target genes in the network.
        n_matched_genes (int)   — network genes present in expression_df.
        coverage_pct    (float) — % of network genes matched; None if unknown.
        resource        (str)
        organism        (str)
        method          (str)
        warnings        (list[str])

    Raises
    ------
    ValueError  if resource or method is not in the supported set.
    """
    import decoupler as dc

    # A caller-supplied network (``_network``) is a custom signature — e.g. a
    # Loveless-derived cell-state signature scored against a bulk cohort
    # (ADR-0006 Role 2). In that case ``resource`` is just a free-form label for
    # outputs/diagnostics, so the built-in-resource gate does not apply.
    if _network is None and resource not in SUPPORTED_RESOURCES:
        raise ValueError(
            f"resource must be one of {sorted(SUPPORTED_RESOURCES)}, got '{resource}'"
        )
    if method not in SUPPORTED_METHODS:
        raise ValueError(
            f"method must be one of {sorted(SUPPORTED_METHODS)}, got '{method}'"
        )

    run_warnings: list[str] = []

    # ── Load network ─────────────────────────────────────────────────────
    if _network is not None:
        net = _network.copy()
    elif resource == "progeny":
        net = dc.op.progeny(organism=organism)
    elif resource == "collectri":
        net = dc.op.collectri(organism=organism)
    else:  # hallmark
        net = dc.op.hallmark(organism=organism)

    # ── Gene coverage ────────────────────────────────────────────────────
    target_col = _find_network_target_col(net)
    expr_genes = set(expression_df.columns.astype(str))

    if target_col:
        network_genes: set[str] = set(net[target_col].astype(str).unique())
        n_network = len(network_genes)
        n_matched = len(network_genes & expr_genes)
        coverage_pct: float | None = (
            round(100 * n_matched / n_network, 1) if n_network > 0 else 0.0
        )

        if coverage_pct is not None and coverage_pct < 20:
            run_warnings.append(
                f"Low gene coverage: {n_matched}/{n_network} network genes "
                f"({coverage_pct}%) found in expression matrix. "
                "Check that column names are HGNC gene symbols."
            )
        elif coverage_pct is not None and coverage_pct < 50:
            run_warnings.append(
                f"Moderate gene coverage: {n_matched}/{n_network} network genes "
                f"({coverage_pct}%) found in expression matrix."
            )
    else:
        n_network = len(net)
        n_matched = 0
        coverage_pct = None
        run_warnings.append(
            "Could not identify target gene column in network — "
            "gene coverage check skipped."
        )

    # ── Sample count ─────────────────────────────────────────────────────
    n_samples = len(expression_df)
    if n_samples < min_n:
        run_warnings.append(
            f"Expression matrix has {n_samples} sample(s), below min_n={min_n}. "
            "Activity estimates from very small cohorts should be interpreted "
            "with caution."
        )

    # ── Pre-flight: zero-overlap guard ───────────────────────────────────
    if target_col and n_matched == 0:
        raise ValueError(
            f"Zero gene overlap: none of the {n_network} {resource} network "
            f"genes were found in the expression columns. "
            "Ensure column names are HGNC gene symbols."
        )

    # ── Run scoring ──────────────────────────────────────────────────────
    # Wrap decoupler's AssertionError (too few overlapping targets per source)
    # into a ValueError with a more actionable message.
    try:
        if method == "ulm":
            acts, pvals = dc.mt.ulm(data=expression_df, net=net)
        elif method == "mlm":
            acts, pvals = dc.mt.mlm(data=expression_df, net=net)
        else:  # zscore
            acts, pvals = dc.mt.zscore(data=expression_df, net=net)
    except AssertionError as exc:
        raise ValueError(
            f"decoupleR could not score with resource='{resource}': {exc}. "
            f"Gene coverage: {n_matched}/{n_network} ({coverage_pct}%). "
            "Each source requires ≥5 overlapping target genes by default. "
            "Ensure expression columns are HGNC gene symbols and sufficient "
            "genes are covered."
        ) from exc

    run_warnings.append(
        f"Activity scores computed with method='{method}', "
        f"resource='{resource}' ({_RESOURCE_DESCRIPTIONS.get(resource, '')}), "
        f"organism='{organism}'. "
        "Input is assumed to be pre-normalised expression. "
        "Do not interpret activity scores as log fold-changes."
    )

    return {
        "activity_df": acts,
        "pvalue_df": pvals,
        "n_samples": n_samples,
        "n_activities": int(acts.shape[1]),
        "n_network_genes": n_network,
        "n_matched_genes": n_matched,
        "coverage_pct": coverage_pct,
        "resource": resource,
        "organism": organism,
        "method": method,
        "warnings": run_warnings,
    }