src/workflows/meta_analysis.py

"""
Late-integration (meta-analysis) engine for ADR-0001 Mode B.

This is the *math* half of the Mode-B contract; the *selection* half lives in
``src/core/combine.py`` (the contract module). The two are deliberately split:
the contract stays closed for modification while this module is open for
extension. A new producer tool becomes meta-analyzable purely by emitting the
result envelope (zero edits here); a new combine method registers a
``required_fields`` key (zero edits to the contract).

What this module provides
-------------------------
1. Four combine strategies, each ``register_strategy(...)``-ed into the contract
   at import time so ``select_strategy(available_fields)`` can dispatch on *which
   envelope columns are present*, never on which tool produced the table:

   ===================  ====================  ========  ======================================
   name                 required_fields       priority  method
   ===================  ====================  ========  ======================================
   inverse_variance     {stat, se}            2         fixed-effect effect-size meta-analysis
   stouffer             {score}               1         Stouffer's Z combination of z-like scores
   fisher               {pvalue}              0         Fisher's combined probability test
   rank_aggregation     {score}               0         robust rank aggregation (Kolde et al.)
   ===================  ====================  ========  ======================================

   Priorities make ``select_strategy`` prefer inverse_variance > stouffer >
   fisher when the available fields allow more than one. ``rank_aggregation`` is
   keyed on ``score`` too but sits *below* stouffer, so it is never auto-selected
   over Stouffer for a score-only table — it is an opt-in alternative (pass
   ``strategy="rank_aggregation"``). (The envelope value-field vocabulary is
   fixed — see ``ENVELOPE_VALUE_FIELDS`` — so a bespoke ``"rank"`` field cannot be
   registered; ranks are derived from ``score`` instead.)

2. ``envelope_from_activity_comparison`` — the adapter that maps
   ``compare_activity_by_group`` output (the real Mode-B source, D4) onto the
   envelope, including the **SE backfill** that makes inverse-variance available.

3. ``cochran_q`` / Cochran's Q + I^2 cross-dataset heterogeneity, reported per
   feature alongside the combined estimate.

4. ``combine_envelopes`` — the top-level entry point: align features across >=2
   envelope tables, pick the strategy from the available fields, run it per
   feature, attach heterogeneity, and BH-adjust the combined p-values.

Strategy-function calling convention
------------------------------------
Every registered strategy ``fn`` takes a dict ``{field: matrix}`` where each
``matrix`` is a ``(n_features, k)`` ``float`` array aligned across the ``k`` input
tables (NaN marks a feature absent from / unusable in that dataset), and returns a
dict of length-``n_features`` arrays: ``combined_score``, ``combined_stat``,
``combined_se``, ``combined_pvalue`` (NaN where a quantity is not defined for that
method). The strategies are intended to be driven through ``combine_envelopes``,
which builds the aligned matrices; the contract module only ever *selects* a
strategy, it never calls ``fn``.

Dependencies: numpy / pandas / scipy only (no new install surface, per the ADR).

See docs/adr/ADR-0001-cross-dataset-integration.md ("Mode B contract") and
docs/adr/ADR-0001-implementation-plan.md (D4).
"""

from __future__ import annotations

from typing import Any

import numpy as np
import pandas as pd
from scipy import stats

from src.core.combine import (
    ENVELOPE_FEATURE_FIELD,
    ENVELOPE_VALUE_FIELDS,
    available_strategies,
    register_strategy,
    select_strategy,
)


# ── Small numeric helpers ────────────────────────────────────────────────
def _two_sided_p_from_z(z: np.ndarray) -> np.ndarray:
    """Two-sided p-value from a z-like statistic (NaN-safe)."""
    return 2.0 * stats.norm.sf(np.abs(z))


def benjamini_hochberg(pvalues: "np.ndarray | list[float]") -> np.ndarray:
    """Benjamini-Hochberg FDR adjustment, NaN-aware.

    Mirrors ``statsmodels.stats.multitest.multipletests(method="fdr_bh")`` (the
    same correction ``compare_activity_by_group`` uses) but in pure numpy so the
    whole engine stays scipy/numpy-only. NaN p-values are passed through as NaN
    and excluded from the rank count ``m``.
    """
    p = np.asarray(pvalues, dtype=float)
    out = np.full(p.shape, np.nan)
    mask = np.isfinite(p)
    m = int(mask.sum())
    if m == 0:
        return out
    pm = p[mask]
    order = np.argsort(pm)
    ranked = pm[order] * m / (np.arange(m) + 1.0)
    # Enforce monotonicity from the largest p downward, then clip into [0, 1].
    ranked = np.minimum.accumulate(ranked[::-1])[::-1]
    adj = np.empty(m, dtype=float)
    adj[order] = np.clip(ranked, 0.0, 1.0)
    out[mask] = adj
    return out


def cochran_q(effect: np.ndarray, weight: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Row-wise Cochran's Q and I^2 across per-dataset estimates.

    Parameters
    ----------
    effect, weight:
        ``(n_features, k)`` arrays of per-dataset effect estimates and their
        weights. A cell is dropped (treated as absent) where either is non-finite
        or the weight is <= 0, so heterogeneity is computed only over the datasets
        that actually contribute to a feature.

    Returns
    -------
    (Q, I2):
        Length-``n_features`` arrays. ``Q = sum_i w_i (theta_i - theta_bar)^2`` with
        ``theta_bar`` the weighted mean; ``I2 = max(0, (Q - (k-1)) / Q) * 100`` as a
        percentage. Both are NaN where fewer than two datasets contribute (k < 2),
        since heterogeneity is undefined for a single estimate.
    """
    eff = np.asarray(effect, dtype=float)
    w = np.asarray(weight, dtype=float)
    usable = np.isfinite(eff) & np.isfinite(w) & (w > 0)
    w = np.where(usable, w, 0.0)
    eff = np.where(usable, eff, 0.0)

    sw = w.sum(axis=1)
    k = usable.sum(axis=1)
    with np.errstate(invalid="ignore", divide="ignore"):
        mean = np.where(sw > 0, (w * eff).sum(axis=1) / sw, np.nan)
        q = (w * (eff - mean[:, None]) ** 2).sum(axis=1)
        df = (k - 1).astype(float)
        i2 = np.where(q > 0, np.clip((q - df) / q, 0.0, 1.0) * 100.0, 0.0)

    enough = k >= 2
    q = np.where(enough, q, np.nan)
    i2 = np.where(enough, i2, np.nan)
    return q, i2


def _empty(n: int) -> np.ndarray:
    return np.full(n, np.nan)


# ── Combine strategies (whole-matrix, NaN-aware) ─────────────────────────
def _stouffer(matrices: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
    """Stouffer's unweighted Z-combination of z-like ``score`` values.

    Z = sum_i z_i / sqrt(k). Combining k identical signals scales the single-table
    z by sqrt(k), so identical strong signals get *more* significant — the
    canonical Stouffer behaviour.
    """
    z = np.asarray(matrices["score"], dtype=float)
    finite = np.isfinite(z)
    k = finite.sum(axis=1)
    zsum = np.nansum(np.where(finite, z, np.nan), axis=1)
    with np.errstate(invalid="ignore", divide="ignore"):
        combined_z = np.where(k > 0, zsum / np.sqrt(k), np.nan)
    n = z.shape[0]
    return {
        "combined_score": combined_z,
        "combined_stat": _empty(n),
        "combined_se": _empty(n),
        "combined_pvalue": _two_sided_p_from_z(combined_z),
    }


def _inverse_variance(matrices: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
    """Fixed-effect inverse-variance (effect-size) meta-analysis on ``stat`` ± ``se``.

    pooled = sum_i (w_i * theta_i) / sum_i w_i  with  w_i = 1 / se_i^2;
    pooled_se = sqrt(1 / sum_i w_i);  z = pooled / pooled_se.

    ``stat`` carries the effect estimate whose standard error is ``se``. For the
    activity adapter that estimate is Cohen's d with its backfilled SE (D4), so
    this performs a standardized-mean-difference meta-analysis.
    """
    theta = np.asarray(matrices["stat"], dtype=float)
    se = np.asarray(matrices["se"], dtype=float)
    usable = np.isfinite(theta) & np.isfinite(se) & (se > 0)

    w = np.where(usable, 1.0 / np.square(np.where(usable, se, 1.0)), 0.0)
    sw = w.sum(axis=1)
    with np.errstate(invalid="ignore", divide="ignore"):
        pooled = np.where(sw > 0, (w * np.where(usable, theta, 0.0)).sum(axis=1) / sw, np.nan)
        pooled_se = np.where(sw > 0, np.sqrt(1.0 / sw), np.nan)
        z = np.where(np.isfinite(pooled_se) & (pooled_se > 0), pooled / pooled_se, np.nan)
    return {
        "combined_score": z,
        "combined_stat": pooled,
        "combined_se": pooled_se,
        "combined_pvalue": _two_sided_p_from_z(z),
    }


def _fisher(matrices: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
    """Fisher's combined probability test on ``pvalue``.

    X = -2 * sum_i ln(p_i) ~ chi^2(2k) under the global null. p-values are clipped
    away from 0 before the log. No signed direction is recoverable from p-values
    alone, so ``combined_score`` is left NaN.
    """
    p = np.asarray(matrices["pvalue"], dtype=float)
    finite = np.isfinite(p)
    k = finite.sum(axis=1)
    pc = np.clip(np.where(finite, p, np.nan), 1e-300, 1.0)
    chi = -2.0 * np.nansum(np.where(finite, np.log(pc), np.nan), axis=1)
    with np.errstate(invalid="ignore"):
        combined_p = np.where(k > 0, stats.chi2.sf(chi, np.maximum(2 * k, 1)), np.nan)
    n = p.shape[0]
    return {
        "combined_score": _empty(n),
        "combined_stat": chi,
        "combined_se": _empty(n),
        "combined_pvalue": combined_p,
    }


def _rank_aggregation(matrices: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
    """Robust rank aggregation (Kolde et al., 2012) over ``score``.

    Within each dataset, features are ranked by descending ``|score|`` and the
    ranks normalised to (0, 1]. For a feature, the sorted normalised ranks
    r_(1) <= ... <= r_(k) give rho = min(1, k * min_i Beta.cdf(r_(i); i, k-i+1)) —
    a p-value for "this feature ranks near the top more consistently than chance".
    ``combined_score = 1 - rho`` (higher = stronger). ``rho`` is returned as the
    combined p-value.

    There is no ``rank`` envelope field (the value-field vocabulary is fixed), so
    this is keyed on ``score`` with priority below Stouffer: it never wins
    auto-dispatch for a score-only table and is reached only on explicit request.
    """
    s = np.asarray(matrices["score"], dtype=float)
    n_features, k = s.shape

    # Normalised within-dataset ranks (1 = strongest |score|), NaN where absent.
    norm = np.full_like(s, np.nan, dtype=float)
    for j in range(k):
        col = s[:, j]
        present = np.where(np.isfinite(col))[0]
        m = present.size
        if m == 0:
            continue
        order = present[np.argsort(-np.abs(col[present]))]  # strongest first
        ranks = np.empty(m, dtype=float)
        ranks[np.argsort(-np.abs(col[present]))] = np.arange(1, m + 1)
        norm[present, j] = ranks / m

    rho = np.full(n_features, np.nan)
    for i in range(n_features):
        r = np.sort(norm[i, np.isfinite(norm[i])])
        kk = r.size
        if kk == 0:
            continue
        idx = np.arange(1, kk + 1)
        beta_cdf = stats.beta.cdf(r, idx, kk - idx + 1)
        rho[i] = min(1.0, float(np.min(beta_cdf)) * kk)

    return {
        "combined_score": 1.0 - rho,
        "combined_stat": _empty(n_features),
        "combined_se": _empty(n_features),
        "combined_pvalue": rho,
    }


# Register at import time — the contract module owns *selection*, this module owns
# the implementations. Re-importing re-registers (idempotent by name).
register_strategy(
    "inverse_variance", {"stat", "se"}, _inverse_variance, priority=2,
    description="Fixed-effect inverse-variance (effect-size) meta-analysis on stat +/- se.",
)
register_strategy(
    "stouffer", {"score"}, _stouffer, priority=1,
    description="Stouffer's Z combination of z-like scores.",
)
register_strategy(
    "fisher", {"pvalue"}, _fisher, priority=0,
    description="Fisher's combined probability test on p-values.",
)
register_strategy(
    "rank_aggregation", {"score"}, _rank_aggregation, priority=0,
    description="Robust rank aggregation (Kolde et al.) on |score|; opt-in alternative to Stouffer.",
)


# ── Envelope adapter: compare_activity_by_group -> envelope (D4) ──────────
_ACTIVITY_REQUIRED = ("effect_size", "statistic", "pvalue", "padj", "n_test", "n_control")


def standard_error_of_cohens_d(d: np.ndarray, n_test: np.ndarray, n_control: np.ndarray) -> np.ndarray:
    """Large-sample standard error of Cohen's d (D4 SE backfill).

    se_d = sqrt((n1 + n2) / (n1 * n2) + d^2 / (2 * (n1 + n2)))

    Always positive for n1, n2 >= 1, so it makes the inverse-variance strategy
    available for activity tables that otherwise carry no ``se`` column.
    """
    d = np.asarray(d, dtype=float)
    n1 = np.asarray(n_test, dtype=float)
    n2 = np.asarray(n_control, dtype=float)
    return np.sqrt((n1 + n2) / (n1 * n2) + d**2 / (2.0 * (n1 + n2)))


def envelope_from_activity_comparison(
    comparison: "dict | pd.DataFrame",
    *,
    dataset_id: str | None = None,
    contrast: str | None = None,
    result_type: str | None = None,
    method: str | None = None,
) -> pd.DataFrame:
    """Map ``compare_activity_by_group`` output onto the Mode-B result envelope.

    Accepts either the full result dict (uses its ``"dataframe"`` key) or the
    dataframe directly. The dataframe is indexed by activity name with columns
    ``effect_size, statistic, pvalue, padj, n_test, n_control`` (see
    ``src/workflows/activity_stats.py``).

    Envelope mapping
    ----------------
    ====================  ===================================================
    envelope column       source
    ====================  ===================================================
    feature               index ("activity")
    score                 effect_size  (Cohen's d — signed, z-like)
    stat                  effect_size  (the effect estimate paired with ``se``)
    se                    DERIVED — standard error of Cohen's d (D4 backfill)
    pvalue                pvalue
    padj                  padj
    n                     n_test + n_control
    ====================  ===================================================

    Deviation from the literal task note (documented intentionally): the prompt
    says ``stat <- statistic`` (the Welch t), but D4 specifies that inverse-variance
    "combines Cohen's d with its backfilled SE". Since the ``inverse_variance``
    strategy reads exactly its required fields ``{stat, se}``, ``stat`` must be the
    effect estimate whose SE is ``se`` — i.e. Cohen's d — for the meta-analysis to
    be coherent (pairing a t-statistic with the SE of d is dimensionally wrong).
    The Welch t is therefore not carried: it is redundant with the standardized
    effect for combination purposes, and ``pvalue`` already encodes its test.

    The optional ``dataset_id`` / ``contrast`` / ``result_type`` / ``method`` tags are
    stamped onto every row when provided, for the result-level compatibility check
    a later caller (the ``decoupler_meta_analyze`` tool) performs.
    """
    df = comparison["dataframe"] if isinstance(comparison, dict) else comparison
    missing = [c for c in _ACTIVITY_REQUIRED if c not in df.columns]
    if missing:
        raise ValueError(
            f"activity comparison is missing column(s) {missing}; expected the "
            f"output of compare_activity_by_group with columns {list(_ACTIVITY_REQUIRED)}"
        )

    d = df["effect_size"].to_numpy(dtype=float)
    n1 = df["n_test"].to_numpy(dtype=float)
    n2 = df["n_control"].to_numpy(dtype=float)
    se = standard_error_of_cohens_d(d, n1, n2)

    envelope = pd.DataFrame(
        {
            ENVELOPE_FEATURE_FIELD: df.index.to_numpy().astype(str),
            "score": d,
            "stat": d,  # effect estimate paired with `se` (see docstring / D4)
            "se": se,
            "pvalue": df["pvalue"].to_numpy(dtype=float),
            "padj": df["padj"].to_numpy(dtype=float),
            "n": n1 + n2,
        }
    )

    tags = {
        "result_type": result_type,
        "method": method,
        "contrast": contrast,
        "dataset_id": dataset_id,
    }
    for tag, value in tags.items():
        if value is not None:
            envelope[tag] = value

    return envelope.reset_index(drop=True)


# ── Feature alignment helpers ────────────────────────────────────────────
def _index_by_feature(table: pd.DataFrame, feature_field: str) -> pd.DataFrame:
    """Return a copy of ``table`` indexed by its (stringified) feature labels."""
    if feature_field in table.columns:
        indexed = table.set_index(feature_field)
    else:
        indexed = table.copy()
    indexed = indexed.copy()
    indexed.index = indexed.index.astype(str)
    indexed.index.name = feature_field
    return indexed


def _available_value_fields(tables: list[pd.DataFrame]) -> list[str]:
    """Envelope value fields present and not all-null in *every* table.

    This is the coarse, column-level availability that drives strategy
    *selection*; per-feature missingness within an available column is handled by
    the NaN-aware strategy math.
    """
    return [
        f
        for f in ENVELOPE_VALUE_FIELDS
        if all(f in t.columns and bool(t[f].notna().any()) for t in tables)
    ]


def _stack_field(tables: list[pd.DataFrame], features: list[str], field: str) -> np.ndarray:
    """Build a ``(n_features, k)`` matrix for ``field`` aligned to ``features``."""
    columns = []
    for t in tables:
        if field in t.columns:
            columns.append(t[field].reindex(features).to_numpy(dtype=float))
        else:
            columns.append(np.full(len(features), np.nan))
    return np.column_stack(columns)


def _strategy_by_name(name: str):
    for strategy in available_strategies():
        if strategy.name == name:
            return strategy
    raise ValueError(
        f"unknown combine strategy '{name}'; registered: "
        f"{sorted(s.name for s in available_strategies())}"
    )


# ── Top-level: combine >=2 envelope tables ───────────────────────────────
def combine_envelopes(
    tables: list[pd.DataFrame],
    *,
    feature_field: str = ENVELOPE_FEATURE_FIELD,
    strategy: str | None = None,
    min_datasets: int = 2,
) -> pd.DataFrame:
    """Meta-analyze >=2 envelope-conforming result tables.

    Steps: align features across tables (union) -> determine which envelope value
    fields are available across all tables -> ``select_strategy`` from those fields
    (field-based dispatch; no tool identity) -> run the strategy per feature ->
    attach cross-dataset heterogeneity (Cochran's Q / I^2) -> BH-adjust the
    combined p-values.

    Parameters
    ----------
    tables:
        List of >=2 tidy envelope DataFrames (one row per feature). Each must carry
        the feature labels in ``feature_field`` (a column) or in its index, plus at
        least one shared envelope value column.
    feature_field:
        Name of the feature column / index (default ``"feature"``).
    strategy:
        Force a named strategy (e.g. ``"rank_aggregation"``). Default ``None`` =
        auto-select the highest-priority strategy the available fields support.
    min_datasets:
        Minimum number of tables required (default 2).

    Returns
    -------
    pd.DataFrame
        One row per feature, sorted by ``padj``. Columns: ``feature``, ``k`` (datasets
        contributing to that feature), ``strategy``, ``combined_score``,
        ``combined_stat``, ``combined_se``, ``combined_pvalue``, ``padj`` (BH), ``Q``,
        ``I2``. ``df.attrs`` records ``strategy``, ``n_tables`` and ``available_fields``.

    Raises
    ------
    ValueError
        If fewer than ``min_datasets`` tables are given, no envelope value field is
        shared across all tables, or a forced ``strategy`` needs an unavailable field.
    src.core.combine.NoCombineStrategyError
        If no registered strategy matches the available fields.
    """
    if len(tables) < min_datasets:
        raise ValueError(
            f"meta-analysis needs at least {min_datasets} result tables, got {len(tables)}"
        )

    indexed = [_index_by_feature(t, feature_field) for t in tables]

    available = _available_value_fields(indexed)
    if not available:
        raise ValueError(
            "no envelope value field is shared across all tables; expected at least "
            f"one of {ENVELOPE_VALUE_FIELDS} present (and non-null) in every table"
        )

    if strategy is None:
        chosen = select_strategy(available)
    else:
        chosen = _strategy_by_name(strategy)
        unmet = chosen.required_fields - set(available)
        if unmet:
            raise ValueError(
                f"strategy '{strategy}' requires field(s) {sorted(unmet)} that are "
                f"not available across all tables (available: {available})"
            )

    # Union of features (sorted for determinism; final output is sorted by padj).
    features = sorted(set().union(*(t.index for t in indexed)))

    # Matrices needed = the strategy's inputs plus whatever heterogeneity needs.
    needed = set(chosen.required_fields)
    if {"stat", "se"} <= set(available):
        het_effect, het_weight_field = "stat", "se"
        needed |= {"stat", "se"}
    elif "score" in available:
        het_effect, het_weight_field = "score", None
        needed |= {"score"}
    else:
        het_effect, het_weight_field = None, None

    matrices = {f: _stack_field(indexed, features, f) for f in needed}

    result = chosen.fn({f: matrices[f] for f in chosen.required_fields})

    # k = datasets where *all* of the strategy's required fields are finite.
    required_finite = np.stack(
        [np.isfinite(matrices[f]) for f in chosen.required_fields], axis=0
    ).all(axis=0)
    k = required_finite.sum(axis=1)

    # Cross-dataset heterogeneity over the best available per-dataset estimates.
    n_features = len(features)
    if het_effect is None:
        q, i2 = _empty(n_features), _empty(n_features)
    elif het_weight_field is not None:
        with np.errstate(divide="ignore", invalid="ignore"):
            weight = 1.0 / np.square(matrices[het_weight_field])
        q, i2 = cochran_q(matrices[het_effect], weight)
    else:
        # No SE -> unit-weight heuristic on the signed score; still flags a dataset
        # whose estimate is far from the others.
        unit = np.isfinite(matrices[het_effect]).astype(float)
        q, i2 = cochran_q(matrices[het_effect], unit)

    out = pd.DataFrame(
        {
            ENVELOPE_FEATURE_FIELD: features,
            "k": k.astype(int),
            "strategy": chosen.name,
            "combined_score": result["combined_score"],
            "combined_stat": result["combined_stat"],
            "combined_se": result["combined_se"],
            "combined_pvalue": result["combined_pvalue"],
            "Q": q,
            "I2": i2,
        }
    )
    # Drop features no dataset actually contributed to the combination.
    out = out[out["k"] > 0].reset_index(drop=True)
    out["padj"] = benjamini_hochberg(out["combined_pvalue"].to_numpy())

    out = out[
        [
            ENVELOPE_FEATURE_FIELD,
            "k",
            "strategy",
            "combined_score",
            "combined_stat",
            "combined_se",
            "combined_pvalue",
            "padj",
            "Q",
            "I2",
        ]
    ]
    out = out.sort_values("padj", na_position="last").reset_index(drop=True)
    out.attrs["strategy"] = chosen.name
    out.attrs["n_tables"] = len(tables)
    out.attrs["available_fields"] = list(available)
    return out


__all__ = [
    "combine_envelopes",
    "envelope_from_activity_comparison",
    "standard_error_of_cohens_d",
    "cochran_q",
    "benjamini_hochberg",
]