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.

3. Input-scale + reporting helpers
   ---------------------------------
   detect_expression_scale
                        Heuristic: is a matrix on a log scale (what ULM assumes)
                        or on a linear scale (raw TPM/CPM, which distorts ULM)?
   plot_activity_landscape
                        Standard cohort-wide landscape figure (per-sample
                        heatmap + mean±SD bar) written to a declared artifact
                        path so the UI's output sweep can embed it.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd

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

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

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

SUPPORTED_INPUT_SCALES: frozenset[str] = frozenset({"auto", "log", "linear"})

# Above this maximum value a matrix is treated as *linear* expression.
# log2-TPM / log-CPM / log microarray intensity all top out around 16-20, so a
# max above 50 means the values were never log-transformed.
LOG_SCALE_MAX_THRESHOLD: float = 50.0

# PROGENy footprint size, pinned explicitly.  decoupler's own default for
# ``dc.op.progeny(top=...)`` has changed across releases, so leaving it implicit
# makes scores irreproducible across upgrades.  500 is the footprint the PROGENy
# publication and decoupler tutorials use for bulk data.
PROGENY_TOP: int = 500

# Maximum number of activity columns drawn in the landscape heatmap.  CollecTRI
# scores ~1 000 TFs, which is unreadable; the most variable ones are kept.
LANDSCAPE_MAX_FEATURES: int = 50

_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 detect_expression_scale(expression_df: pd.DataFrame) -> dict[str, Any]:
    """
    Heuristically decide whether an expression matrix is log-scale or linear.

    ULM/MLM (and PROGENy's weights) assume roughly symmetric, log-scale input.
    Linear TPM/CPM is dominated by a handful of very highly expressed genes, so
    feeding it to ULM distorts every activity score.  This detector is the guard
    on that: it is deliberately conservative and looks only at the value range.

    Returns
    -------
    dict with keys:
        scale        (str)   — "log", "linear", or "unknown" (all-NaN input).
        max_value    (float | None)
        min_value    (float | None)
        has_negative (bool)
        skew         (float | None) — Fisher-Pearson skewness, diagnostic only.
    """
    values = expression_df.to_numpy(dtype=float, copy=False)
    finite = values[np.isfinite(values)]

    if finite.size == 0:
        return {
            "scale": "unknown",
            "max_value": None,
            "min_value": None,
            "has_negative": False,
            "skew": None,
        }

    max_value = float(np.max(finite))
    min_value = float(np.min(finite))
    has_negative = min_value < 0

    std = float(np.std(finite))
    skew = float(np.mean(((finite - np.mean(finite)) / std) ** 3)) if std > 0 else 0.0

    # Negative values mean the matrix is already log-transformed and/or centred
    # (log-ratio, VST, z-scored) — linear abundances are non-negative.
    if has_negative:
        scale = "log"
    elif max_value > LOG_SCALE_MAX_THRESHOLD:
        scale = "linear"
    else:
        scale = "log"

    return {
        "scale": scale,
        "max_value": max_value,
        "min_value": min_value,
        "has_negative": has_negative,
        "skew": skew,
    }


def plot_activity_landscape(
    activity_df: pd.DataFrame,
    out_path: str | Path,
    title: str = "Activity landscape",
    max_features: int = LANDSCAPE_MAX_FEATURES,
) -> dict[str, Any]:
    """
    Render the standard cohort-wide activity landscape figure.

    Two stacked panels:
      1. samples × activities heatmap (per-sample scores, the "landscape"),
      2. cohort mean ± SD bar chart per activity.

    Written by the tool itself so the figure always lands in ``OUTPUT_DIR`` and
    is returned as a declared artifact — agent-authored matplotlib lands wherever
    the agent happens to choose and is never picked up by the UI's output sweep.

    When there are more than ``max_features`` activities, the most variable
    ``max_features`` are plotted (reported in the return dict).

    Returns
    -------
    dict with keys: path, n_samples, n_features_plotted, n_features_total,
    features_plotted, truncated.
    """
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    out_path = Path(out_path)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    n_total = int(activity_df.shape[1])
    plot_df = activity_df
    truncated = n_total > max_features
    if truncated:
        keep = activity_df.std(axis=0).sort_values(ascending=False).head(max_features).index
        plot_df = activity_df.loc[:, keep]

    # Stable, readable ordering: highest mean activity on the left.
    order = plot_df.mean(axis=0).sort_values(ascending=False).index
    plot_df = plot_df.loc[:, order]

    n_features = int(plot_df.shape[1])
    n_samples = int(plot_df.shape[0])

    width = max(6.0, min(0.32 * n_features + 3.0, 24.0))
    height = max(6.0, min(0.06 * n_samples + 5.0, 20.0))

    fig, (ax_heat, ax_bar) = plt.subplots(
        2, 1, figsize=(width, height), height_ratios=[3, 1], constrained_layout=True
    )

    vmax = float(np.nanmax(np.abs(plot_df.to_numpy(dtype=float)))) if n_features else 1.0
    vmax = vmax if vmax > 0 else 1.0
    im = ax_heat.imshow(
        plot_df.to_numpy(dtype=float),
        aspect="auto",
        cmap="RdBu_r",
        vmin=-vmax,
        vmax=vmax,
        interpolation="nearest",
    )
    ax_heat.set_xticks(range(n_features))
    ax_heat.set_xticklabels(plot_df.columns, rotation=90, fontsize=7)
    ax_heat.set_ylabel(f"Samples (n={n_samples})")
    # Sample labels are unreadable past a few dozen rows; the heatmap is a
    # cohort-level view, and per-sample values live in the CSV.
    if n_samples <= 40:
        ax_heat.set_yticks(range(n_samples))
        ax_heat.set_yticklabels(plot_df.index, fontsize=6)
    else:
        ax_heat.set_yticks([])
    ax_heat.set_title(title)
    fig.colorbar(im, ax=ax_heat, label="Activity score", fraction=0.02, pad=0.01)

    means = plot_df.mean(axis=0)
    sds = plot_df.std(axis=0)
    ax_bar.bar(range(n_features), means.to_numpy(), yerr=sds.to_numpy(), capsize=2, color="#4C72B0")
    ax_bar.axhline(0, color="black", linewidth=0.8)
    ax_bar.set_xticks(range(n_features))
    ax_bar.set_xticklabels(plot_df.columns, rotation=90, fontsize=7)
    ax_bar.set_ylabel("Mean ± SD")

    fig.savefig(out_path, dpi=200, bbox_inches="tight")
    plt.close(fig)

    return {
        "path": str(out_path.resolve()),
        "n_samples": n_samples,
        "n_features_plotted": n_features,
        "n_features_total": n_total,
        "features_plotted": list(plot_df.columns),
        "truncated": truncated,
    }


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,
    input_scale: str = "auto",
    progeny_top: int = PROGENY_TOP,
    _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.  The scale is checked rather than
        assumed: with ``input_scale="auto"`` (the default) a matrix that looks
        linear (e.g. raw TPM) is log2(x+1)-transformed before scoring, and the
        transform actually applied is reported in the return dict.

    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.
    input_scale:
        "auto"   — detect the scale (see detect_expression_scale) and log2(x+1)
                   transform the matrix when it looks linear, with a warning.
        "log"    — trust the caller: never transform.  A matrix that still looks
                   linear raises a warning (it is very likely a mistake).
        "linear" — always log2(x+1) transform.
    progeny_top:
        PROGENy footprint size (genes per pathway) passed to dc.op.progeny.
        Pinned by default (PROGENY_TOP) so scores are reproducible across
        decoupler upgrades; ignored for other resources.
    _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)
        input_scale_requested (str)   — the input_scale argument.
        input_scale_detected  (str)   — "log", "linear", or "unknown".
        applied_transform     (str)   — "none" or "log2(x+1)".
        expression_max        (float | None) — max input value, pre-transform.
        network_top     (int | None)  — pinned PROGENy footprint size, else None.
        warnings        (list[str])

    Raises
    ------
    ValueError  if resource, method, or input_scale is not in the supported set,
                or if a log transform is requested on a matrix with negatives.
    """
    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}'")
    if input_scale not in SUPPORTED_INPUT_SCALES:
        raise ValueError(
            f"input_scale must be one of {sorted(SUPPORTED_INPUT_SCALES)}, got '{input_scale}'"
        )

    run_warnings: list[str] = []

    # ── Input scale: ULM assumes log-scale, symmetric values ─────────────
    # Linear TPM/CPM is dominated by a few very highly expressed genes, which
    # silently distorts every activity score, so the scale is checked (and by
    # default corrected) rather than assumed from the docstring.
    scale_info = detect_expression_scale(expression_df)
    detected = scale_info["scale"]
    applied_transform = "none"

    if input_scale == "linear" or (input_scale == "auto" and detected == "linear"):
        if scale_info["has_negative"]:
            raise ValueError(
                "Cannot log-transform an expression matrix containing negative "
                "values. Negative values indicate the data is already log-scale "
                "or centred — pass input_scale='log'."
            )
        expression_df = np.log2(expression_df + 1)
        applied_transform = "log2(x+1)"
        run_warnings.append(
            f"Input looked linear (max={scale_info['max_value']:.1f} > "
            f"{LOG_SCALE_MAX_THRESHOLD:g}); applied a log2(x+1) transform before "
            "scoring. ULM/PROGENy assume roughly symmetric log-scale input — "
            "scoring linear TPM/CPM lets a few high-expression genes dominate. "
            "Pass input_scale='log' to suppress this if the data really is log-scale."
        )
    elif input_scale == "log" and detected == "linear":
        run_warnings.append(
            f"input_scale='log' was requested but the matrix looks linear "
            f"(max={scale_info['max_value']:.1f} > {LOG_SCALE_MAX_THRESHOLD:g}). "
            "No transform was applied, so activity scores may be distorted by "
            "high-expression genes. Verify the dataset's data_level."
        )

    # ── Load network ─────────────────────────────────────────────────────
    network_top: int | None = None
    if _network is not None:
        net = _network.copy()
    elif resource == "progeny":
        # ``top`` is pinned explicitly: decoupler's default has shifted across
        # releases, so an implicit call makes scores irreproducible on upgrade.
        network_top = progeny_top
        net = dc.op.progeny(organism=organism, top=progeny_top)
    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}'. "
        f"Input scale detected='{detected}', transform applied='{applied_transform}'"
        + (f", PROGENy top={network_top}" if network_top is not None else "")
        + ". 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,
        "input_scale_requested": input_scale,
        "input_scale_detected": detected,
        "applied_transform": applied_transform,
        "expression_max": scale_info["max_value"],
        "network_top": network_top,
        "warnings": run_warnings,
    }