src/workflows/pseudobulk.py

"""
Pseudobulk aggregation — the required Path P step before any bulk DE.

Why this exists
---------------
Single-cell datasets (`modality: sc_rnaseq`, ADR-0006 / Path P) must NEVER be
handed to DESeq2 or limma per cell.  Cells from one donor are not independent
replicates: treating them as such inflates n from ~dozens of samples to ~tens of
thousands of cells and produces p-values that are essentially all significant.
The standard fix is to **sum raw counts within each biological sample** first,
recovering a genuine sample x genes counts matrix with one row per donor, which
then takes the ordinary bulk Path A (DESeq2) route.

This module is the aggregation half of that.  It is deliberately generic — it
knows nothing about any specific dataset, only about obs columns the caller
names.

Design notes
------------
- **Sum, not mean.** DESeq2 models counts and their mean-variance relationship;
  averaging destroys the count scale and the library-size information its size
  factors depend on.  ``mode="mean"`` is offered for scoring-style uses but the
  tool warns that the result is not DESeq2 input.
- **Raw counts in, or nothing.** Aggregating already-normalised values sums
  per-cell-normalised numbers, which is meaningless.  Guarded, not assumed.
- **Small groups are dropped, loudly.** A "sample" built from 3 cells is noise;
  it is excluded and reported rather than silently carried into DE.
- Aggregation itself delegates to ``decoupler.pp.pseudobulk`` so the numerics
  match the rest of the decoupleR pipeline; everything around it (validation,
  filtering, diagnostics) is ours.
"""

from __future__ import annotations

from typing import Any

import numpy as np
import pandas as pd

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

SUPPORTED_MODES: frozenset[str] = frozenset({"sum", "mean"})

# A pseudobulk sample built from very few cells is dominated by sampling noise.
# 10 is the common floor in the pseudobulk literature and decoupler's tutorials.
DEFAULT_MIN_CELLS: int = 10

# Total summed counts below this usually means a failed/empty library.
DEFAULT_MIN_COUNTS: int = 1_000

# Below this many pseudobulk samples, a DE contrast has no usable replication.
MIN_SAMPLES_FOR_DE: int = 4


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _matrix_values(adata, layer: str | None) -> np.ndarray:
    """Return a bounded 1-D sample of the matrix values, dense, for inspection."""
    import scipy.sparse as sp

    X = adata.layers[layer] if layer else adata.X
    data = X.data if sp.issparse(X) else np.asarray(X).ravel()
    return np.asarray(data[:100_000], dtype=float)


def looks_like_raw_counts(adata, layer: str | None = None) -> bool:
    """Heuristic: raw UMI counts are non-negative and integer-valued."""
    sample = _matrix_values(adata, layer)
    if sample.size == 0:
        return False
    if (sample < 0).any():
        return False
    return bool(np.allclose(sample, np.round(sample)))


def _total_counts_per_row(adata) -> np.ndarray:
    import scipy.sparse as sp

    X = adata.X
    totals = X.sum(axis=1) if sp.issparse(X) else np.asarray(X).sum(axis=1)
    return np.asarray(totals).ravel()


# ---------------------------------------------------------------------------
# Aggregation
# ---------------------------------------------------------------------------


def aggregate_to_pseudobulk(
    adata,
    sample_col: str,
    groups_col: str | None = None,
    layer: str | None = None,
    mode: str = "sum",
    min_cells: int = DEFAULT_MIN_CELLS,
    min_counts: int = DEFAULT_MIN_COUNTS,
    require_raw_counts: bool = True,
    filter_genes: bool = True,
    filter_group: str | None = None,
) -> dict[str, Any]:
    """
    Aggregate a per-cell AnnData into a pseudobulk samples x genes AnnData.

    Parameters
    ----------
    adata:
        Per-cell AnnData.  ``adata.X`` (or ``layer``) must hold RAW counts.
    sample_col:
        obs column identifying the biological sample / donor / library.  This
        is the unit of replication — one pseudobulk row per value (per group,
        when ``groups_col`` is given).
    groups_col:
        Optional obs column (e.g. cell type / cluster).  When given, aggregation
        is per sample x group, so a contrast can be run within one cell type.
    layer:
        Optional layer name holding raw counts, when ``X`` has been normalised.
    mode:
        "sum" (default; the only correct input for DESeq2) or "mean".
    min_cells:
        Drop pseudobulk samples aggregated from fewer than this many cells.
    min_counts:
        Drop pseudobulk samples whose total counts fall below this.
    require_raw_counts:
        When True (default), raise if the matrix does not look like raw counts.
    filter_genes:
        When True (default), drop lowly-expressed genes after aggregation via
        ``dc.pp.filter_by_expr`` (decoupler's port of edgeR ``filterByExpr``).
        A per-cell matrix is extremely sparse, so an unfiltered pseudobulk object
        carries thousands of all-zero and near-zero genes.  Leaving them in makes
        DESeq2's dispersion trend fail to converge (it falls back to a mean-based
        trend), inflates the multiple-testing burden with untestable genes, and
        manufactures implausible log2FCs off near-zero group means.  Verified on
        the real GSE155698 subset — see the module tests.
    filter_group:
        obs column passed to ``filter_by_expr`` as the grouping (normally the DE
        design factor), so a gene expressed in only one condition is kept.

    Returns
    -------
    dict with keys:
        adata_pseudobulk (AnnData)  — samples x genes, `psbulk_n_cells` in obs.
        n_cells_in       (int)
        n_samples_out    (int)
        n_genes          (int)
        sample_col / groups_col / mode / layer
        dropped_low_cells  (list[dict])  — group + n_cells for each drop.
        dropped_low_counts (list[dict])
        cells_per_sample   (dict)        — retained pseudobulk sample -> n cells.
        warnings           (list[str])

    Raises
    ------
    ValueError  on unknown mode/columns, non-count input, or nothing surviving.
    """
    import decoupler as dc

    run_warnings: list[str] = []

    # ── Validate arguments ────────────────────────────────────────────────
    if mode not in SUPPORTED_MODES:
        raise ValueError(f"mode must be one of {sorted(SUPPORTED_MODES)}, got '{mode}'")
    if sample_col not in adata.obs.columns:
        raise ValueError(
            f"sample_col '{sample_col}' not found in obs. "
            f"Available columns: {list(adata.obs.columns)}"
        )
    if groups_col is not None and groups_col not in adata.obs.columns:
        raise ValueError(
            f"groups_col '{groups_col}' not found in obs. "
            f"Available columns: {list(adata.obs.columns)}"
        )
    if layer is not None and layer not in adata.layers:
        raise ValueError(
            f"layer '{layer}' not found. Available layers: {list(adata.layers.keys())}"
        )

    # ── Guard: raw counts in, or nothing ──────────────────────────────────
    # Summing per-cell-normalised values is meaningless, and the resulting
    # matrix would then sail through DESeq2's own integer check having lost the
    # library-size information its size factors depend on.
    is_counts = looks_like_raw_counts(adata, layer)
    if not is_counts:
        msg = (
            "Matrix does not look like raw counts (values are negative and/or "
            "non-integer). Pseudobulk aggregation must run on RAW counts — "
            "summing normalised per-cell values is not meaningful, and DESeq2 "
            "downstream requires true counts. Pass layer='counts' (or whichever "
            "layer holds them)."
        )
        if require_raw_counts:
            raise ValueError(msg)
        run_warnings.append(msg)

    if mode == "mean":
        run_warnings.append(
            "mode='mean' does not produce DESeq2 input — averaging destroys the "
            "count scale and library-size information. Use mode='sum' for any "
            "differential-expression contrast."
        )

    n_cells_in = int(adata.n_obs)

    # ── Aggregate ─────────────────────────────────────────────────────────
    pdata = dc.pp.pseudobulk(
        adata,
        sample_col=sample_col,
        groups_col=groups_col,
        layer=layer,
        mode=mode,
        skip_checks=True,  # our own guards above are stricter and better-messaged
    )

    # ── Cell counts per pseudobulk sample ─────────────────────────────────
    # decoupler records this, but the key has moved across versions — recompute
    # from the source obs so the number is ours and always present.
    keys = [sample_col] + ([groups_col] if groups_col else [])
    counts = adata.obs.groupby(keys, observed=True).size()

    def _n_cells_for(obs_row) -> int:
        key = obs_row[sample_col] if not groups_col else (obs_row[sample_col], obs_row[groups_col])
        try:
            return int(counts.get(key, 0))
        except TypeError:
            return 0

    pdata.obs["psbulk_n_cells"] = [_n_cells_for(r) for _, r in pdata.obs.iterrows()]
    pdata.obs["psbulk_total_counts"] = _total_counts_per_row(pdata)

    # ── Filter undersized pseudobulk samples ──────────────────────────────
    dropped_low_cells = [
        {"sample": str(idx), "n_cells": int(row["psbulk_n_cells"])}
        for idx, row in pdata.obs.iterrows()
        if row["psbulk_n_cells"] < min_cells
    ]
    keep_cells = pdata.obs["psbulk_n_cells"] >= min_cells
    pdata = pdata[keep_cells].copy()

    # ``min_counts`` is a library-size threshold and only means anything on the
    # count scale. Under mode="mean" the row totals are per-cell averages, which
    # are ~3 orders of magnitude smaller, so applying it there would silently
    # delete every sample. Skip it rather than filter on a meaningless number.
    dropped_low_counts: list[dict] = []
    if mode == "sum":
        dropped_low_counts = [
            {"sample": str(idx), "total_counts": float(row["psbulk_total_counts"])}
            for idx, row in pdata.obs.iterrows()
            if row["psbulk_total_counts"] < min_counts
        ]
        keep_counts = pdata.obs["psbulk_total_counts"] >= min_counts
        pdata = pdata[keep_counts].copy()
    elif min_counts:
        run_warnings.append(
            f"min_counts={min_counts} was ignored: it is a library-size threshold "
            "and is only meaningful for mode='sum'. Cell-count filtering "
            f"(min_cells={min_cells}) still applied."
        )

    if dropped_low_cells:
        run_warnings.append(
            f"Dropped {len(dropped_low_cells)} pseudobulk sample(s) built from "
            f"fewer than min_cells={min_cells} cells — too few cells to be a "
            "reliable expression profile."
        )
    if dropped_low_counts:
        run_warnings.append(
            f"Dropped {len(dropped_low_counts)} pseudobulk sample(s) with total "
            f"counts below min_counts={min_counts}."
        )

    if pdata.n_obs == 0:
        raise ValueError(
            "No pseudobulk samples survived filtering "
            f"(min_cells={min_cells}, min_counts={min_counts}). "
            f"Started from {n_cells_in} cells across "
            f"{adata.obs[sample_col].nunique()} value(s) of '{sample_col}'. "
            "Lower the thresholds or check that sample_col is the donor/library "
            "column rather than a per-cell identifier."
        )

    # ── Drop lowly-expressed genes ────────────────────────────────────────
    # Single-cell matrices are extremely sparse, so a raw pseudobulk object is
    # mostly untestable genes. On the real GSE155698 subset, 8,821 of 36,601
    # genes were all-zero and 12,404 were detected in fewer than 3 samples;
    # feeding that to DESeq2 made the dispersion trend fail to converge and
    # produced |log2FC| >= 10 artefacts off near-zero group means (which the
    # ADR-0002 sanity layer then correctly flagged as critical).
    n_genes_before = int(pdata.n_vars)
    n_genes_filtered = 0
    if filter_genes and mode == "sum":
        if filter_group is not None and filter_group not in pdata.obs.columns:
            run_warnings.append(
                f"filter_group '{filter_group}' is not an obs column of the "
                "pseudobulk object — gene filtering ran ungrouped."
            )
            filter_group = None
        try:
            genes = dc.pp.filter_by_expr(pdata, group=filter_group, inplace=False)
            if genes is not None:
                kept = np.asarray(genes)
                # decoupler returns either a boolean mask or the kept gene names.
                pdata = pdata[:, kept].copy()
                n_genes_filtered = n_genes_before - int(pdata.n_vars)
        except Exception as exc:  # never fail the aggregation over a filter
            run_warnings.append(f"Gene filtering skipped ({exc}).")

        if n_genes_filtered:
            run_warnings.append(
                f"Dropped {n_genes_filtered} lowly-expressed gene(s) of "
                f"{n_genes_before} (edgeR filterByExpr via dc.pp.filter_by_expr"
                + (f", grouped by '{filter_group}'" if filter_group else "")
                + f"), leaving {pdata.n_vars}. Single-cell data is sparse: keeping "
                "all-zero genes breaks DESeq2's dispersion-trend fit and "
                "manufactures large fold-changes off near-zero means."
            )
    elif filter_genes and mode != "sum":
        run_warnings.append(
            "Gene filtering skipped: filter_by_expr is a count-based filter and "
            "only applies to mode='sum'."
        )

    # ── Replication diagnostics ───────────────────────────────────────────
    if pdata.n_obs < MIN_SAMPLES_FOR_DE:
        run_warnings.append(
            f"Only {pdata.n_obs} pseudobulk sample(s) remain. A DE contrast needs "
            "at least ~2 replicates per group; results from this few samples are "
            "not interpretable."
        )

    # A per-cell identifier passed as sample_col is the classic mistake — it
    # yields ~one cell per "sample" and defeats the whole point of pseudobulk.
    if mode == "sum" and n_cells_in > 0:
        mean_cells = float(np.mean(pdata.obs["psbulk_n_cells"])) if pdata.n_obs else 0.0
        if mean_cells < 2:
            run_warnings.append(
                f"Pseudobulk samples average {mean_cells:.1f} cells each — "
                f"'{sample_col}' looks like a per-cell identifier rather than a "
                "biological sample/donor column. Check the obs columns."
            )

    run_warnings.append(
        f"Aggregated {n_cells_in} cells into {pdata.n_obs} pseudobulk sample(s) "
        f"by '{sample_col}'"
        + (f" x '{groups_col}'" if groups_col else "")
        + f" using mode='{mode}'"
        + (f" on layer '{layer}'" if layer else "")
        + ". Cells within a sample are not independent replicates — the "
        "pseudobulk sample is the unit of replication for downstream DE."
    )

    cells_per_sample = {str(idx): int(row["psbulk_n_cells"]) for idx, row in pdata.obs.iterrows()}

    return {
        "adata_pseudobulk": pdata,
        "n_cells_in": n_cells_in,
        "n_samples_out": int(pdata.n_obs),
        "n_genes": int(pdata.n_vars),
        "n_genes_before_filter": n_genes_before,
        "n_genes_filtered_out": n_genes_filtered,
        "filter_genes": filter_genes,
        "filter_group": filter_group,
        "sample_col": sample_col,
        "groups_col": groups_col,
        "mode": mode,
        "layer": layer,
        "input_looked_like_counts": is_counts,
        "dropped_low_cells": dropped_low_cells,
        "dropped_low_counts": dropped_low_counts,
        "cells_per_sample": cells_per_sample,
        "warnings": run_warnings,
    }


def summarize_design(pdata, design_factor: str | None) -> dict[str, Any]:
    """
    Report per-level replication for a candidate DE design factor.

    Returned separately from aggregation so the agent can check a contrast is
    powered *before* paying for DESeq2.  ``usable`` is False when any level has
    fewer than 2 pseudobulk replicates — DESeq2 cannot estimate dispersion from
    a single sample per group.
    """
    if design_factor is None:
        return {"design_factor": None, "levels": {}, "usable": None, "note": None}
    if design_factor not in pdata.obs.columns:
        return {
            "design_factor": design_factor,
            "levels": {},
            "usable": False,
            "note": (
                f"design_factor '{design_factor}' is not an obs column of the "
                f"pseudobulk object. Available: {list(pdata.obs.columns)}"
            ),
        }

    counts = pdata.obs[design_factor].value_counts(dropna=False)
    levels = {str(k): int(v) for k, v in counts.items()}
    under = [k for k, v in levels.items() if v < 2]
    return {
        "design_factor": design_factor,
        "levels": levels,
        "usable": not under,
        "note": (
            f"Level(s) {under} have fewer than 2 pseudobulk replicates — DESeq2 "
            "cannot estimate dispersion for them. Merge, drop, or choose another "
            "factor."
            if under
            else f"All {len(levels)} level(s) have >= 2 pseudobulk replicates."
        ),
    }


def pseudobulk_metadata_frame(pdata) -> pd.DataFrame:
    """Return the pseudobulk obs as a plain DataFrame for CSV export."""
    return pdata.obs.copy()