src/tools/integration_tools.py

"""
Cross-dataset integration tools (ADR-0001, Layer 1).

This MCP server exposes the cross-dataset result-combination paths:

1. decoupler_meta_analyze: *late integration* — combine >=2 per-dataset result
   files (same result_type + contrast) into one meta-analyzed table with
   per-feature cross-dataset heterogeneity (Cochran's Q / I^2). For INDEPENDENT
   cohorts (plan mode == "late").

2. decoupler_normalization_concordance: *same-cohort sensitivity check* — compare
   >=2 results that are the SAME samples quantified/normalized different ways
   (sibling variants; plan mode == "concordance"). It does NOT combine — it
   reports descriptive agreement (Pearson/Spearman, sign-concordance, significant-
   call overlap), because meta-analyzing identical samples double-counts the
   cohort. Routed here by get_integration_plan's mode == "concordance".

3. decoupler_integrate_datasets: *early integration — DE* — POOL >=2 datasets into
   one matrix (feature intersection + a `batch` obs key) and run a single batch-aware
   DE (`dataset` as a covariate). Only for datasets the plan rules "early" (shared
   feature space + poolable data_level); it re-checks the plan and refuses/reroutes
   otherwise. (ADR-0001 T9, Mode A.)

4. decoupler_pool_cohorts: *early integration — SCORING* — POOL >=2 datasets and
   ComBat-correct (scanpy.pp.combat on the log-normalised matrix) into ONE
   pre-normalised matrix for PER-SAMPLE activity scoring (dataset_score_bulk_samples).
   The no-design-matrix counterpart to decoupler_integrate_datasets: scoring has no
   design matrix to carry a batch covariate, so the pooled matrix is batch-corrected
   before scoring. Same "early"-only plan gate. (ADR-0001 item 10, Mode A.)

Design (ADR-0001 Mode B):
- The *math* is generic over the result envelope and lives in
  src/workflows/meta_analysis.py (``combine_envelopes``); *strategy selection* is
  field-based in src/core/combine.py. This tool is the Layer-1 boundary: it reads
  the producer tools' native CSVs into the envelope, runs the engine, and writes
  the combined table. It never picks the combine math by tool identity.
- Which dataset can be combined with which is decided upstream by
  ``dataset_get_integration_plan`` (mode early/late/refuse); on a "late" verdict
  the specialist runs each dataset's pipeline, then calls this tool.
"""

from __future__ import annotations

from datetime import datetime
from pathlib import Path
from typing import Annotated, Literal

import pandas as pd
from fastmcp import FastMCP

from src.core.combine import (
    CANONICAL_RESULT_TYPES,
    ENVELOPE_FEATURE_FIELD,
    ENVELOPE_VALUE_FIELDS,
)

# Mode-A early integration (T9 + item 10): the plan gate, the combined-AnnData
# builder, the ComBat pre-correction (scoring path), and the batch-aware DE tool.
# Module-level so tests can monkeypatch them.
from src.datasets.registry import get_integration_plan
from src.tools.rna._base import OUTPUT_DIR
from src.tools.rna.analysis import decoupler_differential_expression
from src.workflows.concordance import concordance_metrics
from src.workflows.integration import batch_correct_for_scoring, build_combined_anndata
from src.workflows.meta_analysis import combine_envelopes

integration_mcp = FastMCP(name="integration")

# Enrichment tools serialise a WIDE activity matrix (1 contrast row x N features);
# the DE tool serialises a LONG table (feature rows x stat/pvalue/padj columns).
# These are the only two native shapes produced by the four `@combinable` rna
# tools — the result_type tells the reader which shape to expect.
_ENRICHMENT_RESULT_TYPES = frozenset(
    {
        "tf_enrichment",
        "pathway_enrichment",
        "hallmark_enrichment",
    }
)

# Default feature-label column per result_type (matches the rna tools' combine
# descriptors). Only used for messaging — the envelope feature column is always
# normalised to ENVELOPE_FEATURE_FIELD.
_DEFAULT_FEATURE_FIELD = {
    "de": "gene",
    "tf_enrichment": "TF",
    "pathway_enrichment": "Pathway",
    "hallmark_enrichment": "Gene_Set",
}

_REQUIRED_SPEC_KEYS = ("path", "result_type", "contrast")


def _spec_to_envelope(spec: dict) -> pd.DataFrame:
    """Read one per-dataset result file into a Mode-B result envelope.

    ``spec`` keys: ``path`` (the result CSV), ``result_type`` (one of
    CANONICAL_RESULT_TYPES), ``contrast`` (the comparison label), optional
    ``dataset_id`` (defaults to the file stem) and optional ``padj_path``
    (enrichment only — a companion p-adjusted CSV).

    Two native shapes are handled, chosen by ``result_type``:

    - enrichment (tf/pathway/hallmark): the activities CSV is WIDE — one contrast
      row x N feature columns. The single row is transposed to one row per
      feature with the activity under the envelope ``score`` column. A
      ``padj_path`` companion, if given, contributes the envelope ``padj`` column.
    - ``de``: the results CSV is already LONG — feature rows with envelope-named
      value columns (``stat``/``pvalue``/``padj``); those present are carried
      through.

    The shared tags (``result_type``/``contrast``/``dataset_id``) are stamped on
    every row so the caller's result-level compatibility check and the combined
    output can identify the source.
    """
    path = Path(spec["path"])
    if not path.exists():
        raise FileNotFoundError(f"result file not found: {path}")
    result_type = spec["result_type"]
    contrast = spec["contrast"]
    dataset_id = spec.get("dataset_id") or path.stem

    raw = pd.read_csv(path, index_col=0)
    if raw.shape[0] == 0 or raw.shape[1] == 0:
        raise ValueError(f"result file is empty: {path}")

    if result_type in _ENRICHMENT_RESULT_TYPES:
        # Wide activities: take the (single) contrast row -> score per feature.
        scores = raw.iloc[0]
        env = pd.DataFrame(
            {
                ENVELOPE_FEATURE_FIELD: scores.index.astype(str),
                "score": pd.to_numeric(scores.to_numpy(), errors="coerce"),
            }
        )
        padj_path = spec.get("padj_path")
        if padj_path:
            padj_raw = pd.read_csv(padj_path, index_col=0)
            padj_row = padj_raw.iloc[0]
            padj_map = {str(k): pd.to_numeric(v, errors="coerce") for k, v in padj_row.items()}
            env["padj"] = env[ENVELOPE_FEATURE_FIELD].map(padj_map)
    else:
        # Long table: feature in the index, envelope value columns by name.
        env = raw.copy()
        env.index = env.index.astype(str)
        # reset_index() makes the former index the first column; name it `feature`.
        env = env.reset_index()
        env = env.rename(columns={env.columns[0]: ENVELOPE_FEATURE_FIELD})
        present = [f for f in ENVELOPE_VALUE_FIELDS if f in env.columns]
        if not present:
            raise ValueError(
                f"result file {path} (result_type='{result_type}') has none of the "
                f"envelope value columns {ENVELOPE_VALUE_FIELDS}; columns present: "
                f"{list(raw.columns)}"
            )
        env = env[[ENVELOPE_FEATURE_FIELD, *present]]

    env["result_type"] = result_type
    env["contrast"] = contrast
    env["dataset_id"] = dataset_id
    return env


def _refusal(reason: str, **extra) -> dict:
    """Uniform refuse payload (mirrors the dataset tools' error-dict convention)."""
    return {"error": reason, "refused": True, **extra}


@integration_mcp.tool
def decoupler_meta_analyze(
    results: Annotated[
        list[dict],
        "Two or more per-dataset result specs to meta-analyze (late integration). "
        "Each item is a dict: {'path': <result CSV>, 'result_type': one of "
        "'de'|'tf_enrichment'|'pathway_enrichment'|'hallmark_enrichment', "
        "'contrast': <label, e.g. 'Tumor.vs.Normal'>, 'dataset_id': <optional "
        "label>, 'padj_path': <optional enrichment p-adjusted CSV>}. For DE use "
        "the de_results_path; for enrichment use the activities CSV (output_path). "
        "ALL items MUST share the same result_type AND contrast.",
    ],
    strategy: Annotated[
        str | None,
        "Optional combine strategy override: 'inverse_variance' | 'stouffer' | "
        "'fisher' | 'rank_aggregation'. Default None auto-selects the best strategy "
        "the available envelope fields support (effect-size > Stouffer > Fisher).",
    ] = None,
    out_prefix: Annotated[str | None, "Output file prefix for the combined results CSV."] = None,
) -> dict:
    """
    Meta-analyze (late integration) >=2 per-dataset result tables into one.

    Use this for the "late" branch of a cross-dataset request: after
    dataset_get_integration_plan returns mode=="late", run each dataset's pipeline
    SEPARATELY (same analysis tool, same contrast/method on each), then call this
    tool with the per-dataset result files.

    It enforces the result-level compatibility check (all inputs share
    result_type AND contrast — a mismatch is refused), reads each result into the
    common envelope, combines them with a field-dispatched strategy (no pooling of
    raw expression), and reports per-feature cross-dataset heterogeneity (Cochran's
    Q and I^2). High I^2 (>50%) flags a feature whose effect is inconsistent across
    cohorts.

    Returns a summary dict with the chosen strategy, the combined results CSV path,
    the number of features and significant features, and the top combined rows.
    Returns an {"error": ..., "refused": True} dict (does not raise) when inputs
    are incompatible, so the agent can surface the reason to the user.
    """
    # --- Validate arity + per-item schema ---------------------------------
    if not isinstance(results, list) or len(results) < 2:
        return _refusal(
            f"Meta-analysis needs at least 2 per-dataset results; got "
            f"{len(results) if isinstance(results, list) else 'a non-list'}."
        )
    for i, spec in enumerate(results):
        if not isinstance(spec, dict):
            return _refusal(f"results[{i}] is not a dict: {spec!r}")
        missing = [k for k in _REQUIRED_SPEC_KEYS if not spec.get(k)]
        if missing:
            return _refusal(f"results[{i}] is missing required key(s) {missing}.")
        rt = spec["result_type"]
        if rt not in CANONICAL_RESULT_TYPES:
            return _refusal(
                f"results[{i}] has unknown result_type '{rt}'; expected one of "
                f"{sorted(CANONICAL_RESULT_TYPES)}."
            )

    # --- Result-level compatibility: shared result_type AND contrast ------
    result_types = {s["result_type"] for s in results}
    contrasts = {s["contrast"] for s in results}
    if len(result_types) > 1:
        return _refusal(
            "Cannot meta-analyze across different result types "
            f"{sorted(result_types)} — combine only like-with-like (e.g. all "
            "tf_enrichment). Run the SAME analysis tool on each dataset first.",
            result_types=sorted(result_types),
        )
    if len(contrasts) > 1:
        return _refusal(
            f"Cannot meta-analyze across different contrasts {sorted(contrasts)} — "
            "every dataset must be analyzed with the same contrast before combining.",
            contrasts=sorted(contrasts),
        )

    result_type = next(iter(result_types))
    contrast = next(iter(contrasts))

    # --- Read each result into the envelope -------------------------------
    try:
        envelopes = [_spec_to_envelope(spec) for spec in results]
    except (FileNotFoundError, ValueError) as exc:
        return _refusal(f"Could not read a result into the envelope: {exc}")

    dataset_ids = [spec.get("dataset_id") or Path(spec["path"]).stem for spec in results]

    # --- Combine (field-based strategy dispatch) --------------------------
    try:
        combined = combine_envelopes(envelopes, strategy=strategy)
    except Exception as exc:  # NoCombineStrategyError, ValueError, ...
        return _refusal(f"Meta-analysis could not combine these results: {exc}")

    strategy_used = combined.attrs.get("strategy", strategy or "auto")

    out_prefix = out_prefix or f"meta_{result_type}_{datetime.now():%Y%m%d_%H%M%S}"
    out_path = OUTPUT_DIR / f"{out_prefix}_meta_analysis.csv"
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    combined.to_csv(out_path, index=False)

    n_features = int(len(combined))
    n_sig = int((combined["padj"] < 0.05).sum()) if "padj" in combined else 0

    return {
        "message": (
            f"Meta-analyzed {len(results)} datasets "
            f"({', '.join(dataset_ids)}) for {result_type} / contrast '{contrast}' "
            f"using the '{strategy_used}' strategy. {n_features} features combined, "
            f"{n_sig} significant (padj < 0.05)."
        ),
        "mode": "late",
        "result_type": result_type,
        "contrast": contrast,
        "datasets_combined": dataset_ids,
        "n_datasets": len(results),
        "strategy": strategy_used,
        "available_fields": combined.attrs.get("available_fields", []),
        "n_features": n_features,
        "n_significant": n_sig,
        "combined_results_path": str(out_path.resolve()),
        "output_path": str(out_path.resolve()),
        "top_results": combined.head(10).round(4).to_dict(orient="records"),
        "heterogeneity_note": (
            "Per-feature Cochran's Q and I^2 are in the combined CSV (columns Q, "
            "I2). I^2 > 50% flags a feature whose effect is inconsistent across "
            "cohorts; report it alongside the combined estimate."
        ),
        "artifacts": [
            {
                "description": "Meta-analysis combined results (combined stat/score/p, padj, Q, I2)",
                "path": str(out_path.resolve()),
            }
        ],
    }


@integration_mcp.tool
def decoupler_normalization_concordance(
    results: Annotated[
        list[dict],
        "Two or more SAME-COHORT variant result specs to compare (e.g. the TPM "
        "and TMM quantifications of one cohort). Each item is a dict: {'path': "
        "<result CSV>, 'result_type': one of 'de'|'tf_enrichment'|"
        "'pathway_enrichment'|'hallmark_enrichment', 'contrast': <label>, "
        "'dataset_id': <variant label, e.g. 'gse205154_sears_tpm'>, 'padj_path': "
        "<optional enrichment p-adjusted CSV>}. For DE use the de_results_path; for "
        "enrichment use the activities CSV (output_path). ALL items MUST share the "
        "same result_type AND contrast.",
    ],
    sig_threshold: Annotated[
        float,
        "Adjusted-p cutoff defining a 'significant' call per variant, for the "
        "significant-call overlap (default 0.05).",
    ] = 0.05,
    out_prefix: Annotated[
        str | None, "Output file prefix for the per-feature concordance CSV."
    ] = None,
) -> dict:
    """
    Compare same-cohort variants descriptively (normalization sensitivity check).

    Use this for the "concordance" branch of a cross-dataset request: when
    dataset_get_integration_plan returns mode=="concordance" (the requested
    datasets are sibling quantifications of ONE cohort — same samples, different
    units, e.g. GSE205154 TPM vs TMM). Run each variant's pipeline SEPARATELY
    (same analysis tool, same contrast on each), then call this tool with the
    per-variant result files.

    Do NOT call decoupler_meta_analyze for sibling variants: meta-analysis assumes
    independent cohorts, so combining identical samples double-counts them
    (Stouffer inflates the score by ~sqrt(N); Cochran's Q / I^2 are 0 by
    construction). This tool combines nothing — it reports how much the
    normalization choice changed the answer:

    - pairwise Pearson and Spearman correlation of the per-feature effect,
    - sign-concordance (fraction of shared features pointing the same way),
    - the spread of the effect across variants, and
    - overlap of the significant-call sets (Jaccard + per-variant-only counts),
      when an adjusted p-value is available.

    Returns a summary dict (mode == "concordance") plus a per-feature CSV. Returns
    an {"error": ..., "refused": True} dict (does not raise) when inputs are
    incompatible, so the agent can surface the reason.
    """
    # --- Validate arity + per-item schema (mirrors decoupler_meta_analyze) ---
    if not isinstance(results, list) or len(results) < 2:
        return _refusal(
            f"Concordance needs at least 2 same-cohort variant results; got "
            f"{len(results) if isinstance(results, list) else 'a non-list'}."
        )
    for i, spec in enumerate(results):
        if not isinstance(spec, dict):
            return _refusal(f"results[{i}] is not a dict: {spec!r}")
        missing = [k for k in _REQUIRED_SPEC_KEYS if not spec.get(k)]
        if missing:
            return _refusal(f"results[{i}] is missing required key(s) {missing}.")
        rt = spec["result_type"]
        if rt not in CANONICAL_RESULT_TYPES:
            return _refusal(
                f"results[{i}] has unknown result_type '{rt}'; expected one of "
                f"{sorted(CANONICAL_RESULT_TYPES)}."
            )

    # --- Result-level compatibility: shared result_type AND contrast ---------
    result_types = {s["result_type"] for s in results}
    contrasts = {s["contrast"] for s in results}
    if len(result_types) > 1:
        return _refusal(
            "Cannot compare across different result types "
            f"{sorted(result_types)} — compare like-with-like (e.g. all "
            "tf_enrichment). Run the SAME analysis tool on each variant first.",
            result_types=sorted(result_types),
        )
    if len(contrasts) > 1:
        return _refusal(
            f"Cannot compare across different contrasts {sorted(contrasts)} — every "
            "variant must be analyzed with the same contrast before comparing.",
            contrasts=sorted(contrasts),
        )

    result_type = next(iter(result_types))
    contrast = next(iter(contrasts))
    dataset_ids = [spec.get("dataset_id") or Path(spec["path"]).stem for spec in results]

    # --- Read each result into the envelope ----------------------------------
    try:
        envelopes = [_spec_to_envelope(spec) for spec in results]
    except (FileNotFoundError, ValueError) as exc:
        return _refusal(f"Could not read a result into the envelope: {exc}")

    # --- Descriptive agreement (no combine) ----------------------------------
    try:
        summary, per_feature = concordance_metrics(
            envelopes, dataset_ids=dataset_ids, sig_threshold=sig_threshold
        )
    except ValueError as exc:
        return _refusal(f"Concordance could not compare these results: {exc}")

    out_prefix = out_prefix or f"concordance_{result_type}_{datetime.now():%Y%m%d_%H%M%S}"
    out_path = OUTPUT_DIR / f"{out_prefix}_concordance.csv"
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    per_feature.to_csv(out_path, index=False)

    overlap = summary["significant_overlap"]
    if summary["n_variants"] == 2:
        corr_msg = f"Pearson r={summary['pearson_r']:.4f}, Spearman r={summary['spearman_r']:.4f}"
    else:
        corr_msg = (
            f"min pairwise Pearson r={summary['pearson_min']:.4f}, "
            f"min Spearman r={summary['spearman_min']:.4f}"
        )
    overlap_msg = (
        f"; {overlap['n_significant_in_all']}/{overlap['n_significant_in_any']} "
        f"significant calls agree (Jaccard {overlap['jaccard_significant']:.3f})"
        if overlap.get("available")
        else ""
    )

    return {
        "message": (
            f"Compared {summary['n_variants']} same-cohort variants "
            f"({', '.join(dataset_ids)}) for {result_type} / contrast '{contrast}' "
            f"on the '{summary['effect_field']}' effect over "
            f"{summary['n_features_shared']} shared features: {corr_msg}; "
            f"sign-concordance {summary['sign_concordance']:.3f}{overlap_msg}. "
            f"This is a normalization sensitivity check, NOT a meta-analysis — the "
            f"variants are the same samples, so results are not combined."
        ),
        "mode": "concordance",
        "result_type": result_type,
        "contrast": contrast,
        "variants_compared": dataset_ids,
        "n_variants": summary["n_variants"],
        "n_features_shared": summary["n_features_shared"],
        "n_features_union": summary["n_features_union"],
        "effect_field": summary["effect_field"],
        "pearson_pairwise": summary["pearson_pairwise"],
        "spearman_pairwise": summary["spearman_pairwise"],
        "pearson_min": summary["pearson_min"],
        "spearman_min": summary["spearman_min"],
        "sign_concordance": summary["sign_concordance"],
        "n_sign_flips": summary["n_sign_flips"],
        "sign_flip_features": summary["sign_flip_features"],
        "max_effect_spread": summary["max_effect_spread"],
        "mean_effect_spread": summary["mean_effect_spread"],
        "significant_overlap": overlap,
        "top_divergent_features": summary["top_divergent_features"],
        "concordance_table_path": str(out_path.resolve()),
        "output_path": str(out_path.resolve()),
        "interpretation_note": (
            "High correlation + sign-concordance ~1.0 + high Jaccard means the "
            "normalization choice barely changes the result (a robustness signal). "
            "Report agreement metrics, not a combined score — there is no added "
            "statistical power here because the variants share all samples. "
            "Variant-only significant calls are normalization-sensitive (marginal) "
            "and should not be treated as robust findings."
        ),
        "artifacts": [
            {
                "description": (
                    "Per-feature concordance table (effect per variant, spread, "
                    "sign agreement, per-variant significance)"
                ),
                "path": str(out_path.resolve()),
            }
        ],
    }


@integration_mcp.tool
def decoupler_integrate_datasets(
    dataset_ids: Annotated[
        list[str],
        "Two or more registered dataset IDs to POOL into one matrix and analyze "
        "together (early integration), e.g. ['tcga_paad', 'paca_au_rnaseq'].",
    ],
    design_factor: Annotated[
        str,
        "obs column holding the contrast groups; must exist in EVERY dataset "
        "(e.g. 'tumor_subtype').",
    ],
    test_group: Annotated[str, "Test arm of the contrast (a value of design_factor)."],
    control_group: Annotated[str, "Control / reference arm of the contrast."],
    method: Annotated[
        Literal["auto", "deseq2", "limma"],
        "DE method on the pooled matrix. 'auto' (default) picks deseq2 for pooled "
        "raw counts and limma otherwise, from the integration plan's poolable "
        "data_level. `dataset` is always modelled as a batch covariate.",
    ] = "auto",
    min_shared_features: Annotated[
        int,
        "Minimum shared gene symbols required to pool; below this the datasets "
        "likely do not share a feature space and the request is refused.",
    ] = 200,
    out_prefix: Annotated[str | None, "Output file prefix."] = None,
) -> dict:
    """
    Early integration (Mode A): pool >=2 datasets into one matrix + one batch-aware DE.

    Use this ONLY when dataset_get_integration_plan returns mode=="early": the
    datasets share a poolable feature space AND data_level, so they can be
    concatenated and analyzed jointly with `dataset` as a batch covariate (more
    power than meta-analysis when valid). This tool RE-CHECKS the plan itself and
    refuses (does not pool) when the verdict is not 'early':

    - 'late'        -> run each dataset separately, then decoupler_meta_analyze.
    - 'concordance' -> same-cohort variants -> decoupler_normalization_concordance.
    - 'refuse'      -> incompatible (e.g. confounded design / cross-modality).

    On 'early' it builds the combined AnnData (feature intersection + a `batch`
    obs key = dataset_id), runs decoupler_differential_expression with
    batch_column='batch' (DESeq2 ~batch+factor for raw counts, limma ~batch+group
    otherwise), and returns the pooled DE result plus per-dataset composition.
    Returns an {"error": ..., "refused": True} dict (does not raise) so the agent
    can surface the reason and reroute.

    Pass design_factor + test_group + control_group: they drive the plan's
    confound check AND the pooled contrast.
    """
    if not isinstance(dataset_ids, list) or len(dataset_ids) < 2:
        return _refusal(
            "Early integration needs >=2 dataset_ids; got "
            f"{len(dataset_ids) if isinstance(dataset_ids, list) else 'a non-list'}."
        )

    try:
        plan = get_integration_plan(
            dataset_ids,
            design_factor=design_factor,
            test_group=test_group,
            control_group=control_group,
        )
    except Exception as exc:  # unknown dataset id, etc.
        return _refusal(f"Could not compute the integration plan: {exc}")

    mode = plan.get("mode")
    if mode != "early":
        reason = plan.get("reason", "")
        reroute = {
            "late": "Run each dataset separately, then call decoupler_meta_analyze.",
            "concordance": (
                "These are same-cohort variants — call decoupler_normalization_concordance instead."
            ),
        }.get(mode, "")
        return _refusal(
            (
                f"Datasets are not eligible for early pooling (plan mode='{mode}'). "
                f"{reason} {reroute}"
            ).strip(),
            mode=mode,
            plan_reason=reason,
            refusal_rules_triggered=plan.get("refusal_rules_triggered", []),
        )

    # Derive the DE method from the poolable data_level unless explicitly set.
    if method == "auto":
        method = "deseq2" if plan.get("poolable_data_level") == "raw_counts" else "limma"

    try:
        built = build_combined_anndata(
            dataset_ids,
            design_factor=design_factor,
            batch_key="batch",
            min_shared_features=min_shared_features,
            out_prefix=(f"{out_prefix}_combined" if out_prefix else None),
        )
    except Exception as exc:
        return _refusal(f"Could not build the pooled dataset: {exc}", mode="early")

    try:
        de = decoupler_differential_expression(
            adata_path=built["output_path"],
            design_factor=design_factor,
            contrast=[design_factor, test_group, control_group],
            method=method,
            batch_column="batch",
            out_prefix=out_prefix,
        )
    except Exception as exc:
        return _refusal(f"Pooled DE failed: {exc}", mode="early", combined=built, method=method)

    return {
        "message": (
            f"Early integration of {len(dataset_ids)} datasets "
            f"({', '.join(dataset_ids)}): pooled {built['n_obs']} samples over "
            f"{built['n_vars']} shared genes and ran {de.get('method_used')} with "
            f"dataset modelled as a batch covariate. "
            f"{de.get('n_significant')} significant genes."
        ),
        "mode": "early",
        "plan_reason": plan.get("reason", ""),
        "datasets": list(dataset_ids),
        "batch_key": "batch",
        "per_batch_n": built["per_batch_n"],
        "n_shared_features": built["n_vars"],
        "n_combined_samples": built["n_obs"],
        "method_used": de.get("method_used"),
        "n_significant": de.get("n_significant"),
        "batch_modeled": True,
        "combined_path": built["output_path"],
        "de_results_path": de.get("output_path"),
        "sanity_warnings": de.get("sanity_warnings"),
        "interpretation_note": (
            "Early-integration (pooled) result: the datasets were concatenated and "
            "analyzed jointly with `dataset` as a batch covariate, so the group "
            "effect is estimated adjusting for dataset. Report the datasets combined, "
            "the per-dataset sample counts, and the batch handling. This is more "
            "powerful than meta-analysis, but valid ONLY because the plan verdict was "
            "'early' (shared feature space + poolable data_level)."
        ),
        "artifacts": [
            {"description": "Pooled DE results (CSV)", "path": de.get("output_path")},
            {"description": "Combined AnnData (h5ad)", "path": built["output_path"]},
        ],
    }


@integration_mcp.tool
def decoupler_pool_cohorts(
    dataset_ids: Annotated[
        list[str],
        "Two or more registered dataset IDs to POOL into one batch-corrected matrix "
        "for PER-SAMPLE activity scoring (early integration), e.g. ['tcga_paad', "
        "'paca_au_rnaseq'].",
    ],
    min_shared_features: Annotated[
        int,
        "Minimum shared gene symbols required to pool; below this the datasets "
        "likely do not share a feature space and the request is refused.",
    ] = 200,
    out_prefix: Annotated[str | None, "Output file prefix."] = None,
) -> dict:
    """
    Early integration for SCORING: pool >=2 datasets + ComBat-correct into one matrix.

    Use this when the user wants PER-SAMPLE activity scores (PROGENy / CollecTRI /
    Hallmark via dataset_score_bulk_samples) across two or more cohorts pooled
    together — NOT a differential-expression contrast. It returns ONE pre-normalised,
    batch-corrected samples x genes matrix; feed its output_path straight to
    dataset_score_bulk_samples, so the scores are no longer cohort-confounded.

    This is the no-design-matrix counterpart to decoupler_integrate_datasets:
    - decoupler_integrate_datasets  -> early DE CONTRAST: models `dataset` as a batch
      COVARIATE in the test (the statistically correct route for DE).
    - decoupler_pool_cohorts (this) -> early per-sample SCORING: there is no design
      matrix to carry a covariate, so the pooled matrix is ComBat-corrected
      (scanpy.pp.combat on the log-normalised values) BEFORE scoring.

    Like decoupler_integrate_datasets it RE-CHECKS dataset_get_integration_plan and
    pools ONLY on mode=='early'; otherwise it refuses (does not pool):
    - 'late'        -> score each cohort separately, or decoupler_meta_analyze a DE.
    - 'concordance' -> same-cohort variants -> decoupler_normalization_concordance.
    - 'refuse'      -> incompatible (e.g. cross-modality / no shared feature space).

    ComBat is run WITHOUT a biological covariate (pure dataset-shift removal); if the
    biological groups are confounded with cohort it can also remove real between-cohort
    biology — stated in the returned interpretation_note. Do NOT use this matrix for DE
    testing (ComBat + naive DE inflates false positives — use decoupler_integrate_datasets).

    Returns an {"error": ..., "refused": True} dict (does not raise) so the agent can
    surface the reason and reroute.
    """
    if not isinstance(dataset_ids, list) or len(dataset_ids) < 2:
        return _refusal(
            "Early pooling needs >=2 dataset_ids; got "
            f"{len(dataset_ids) if isinstance(dataset_ids, list) else 'a non-list'}."
        )

    try:
        plan = get_integration_plan(dataset_ids)
    except Exception as exc:  # unknown dataset id, etc.
        return _refusal(f"Could not compute the integration plan: {exc}")

    mode = plan.get("mode")
    if mode != "early":
        reason = plan.get("reason", "")
        reroute = {
            "late": (
                "Score each cohort separately, or run a contrast per cohort and "
                "decoupler_meta_analyze the results."
            ),
            "concordance": (
                "These are same-cohort variants — call decoupler_normalization_concordance instead."
            ),
        }.get(mode, "")
        return _refusal(
            (
                f"Datasets are not eligible for early pooling (plan mode='{mode}'). "
                f"{reason} {reroute}"
            ).strip(),
            mode=mode,
            plan_reason=reason,
            refusal_rules_triggered=plan.get("refusal_rules_triggered", []),
        )

    # Pool into one raw combined matrix (no contrast needed for scoring).
    try:
        built = build_combined_anndata(
            dataset_ids,
            design_factor=None,
            batch_key="batch",
            min_shared_features=min_shared_features,
            out_prefix=(f"{out_prefix}_pooled" if out_prefix else None),
        )
    except Exception as exc:
        return _refusal(f"Could not build the pooled dataset: {exc}", mode="early")

    # ComBat-correct the pooled matrix for per-sample scoring.
    import scanpy as sc

    raw_pooled_path = built["output_path"]
    try:
        pooled = sc.read_h5ad(raw_pooled_path)
        corrected, info = batch_correct_for_scoring(
            pooled, batch_key="batch", data_level=plan.get("poolable_data_level")
        )
    except Exception as exc:
        return _refusal(f"ComBat batch correction failed: {exc}", mode="early", combined=built)

    prefix = out_prefix or ("pooled_" + "_".join(str(d) for d in dataset_ids))[:80]
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    out_path = OUTPUT_DIR / f"{prefix}_combat.h5ad"
    corrected.write_h5ad(out_path)
    # Drop the intermediate UN-corrected pooled matrix so it cannot be scored by mistake.
    Path(raw_pooled_path).unlink(missing_ok=True)

    return {
        "message": (
            f"Early integration (scoring) of {len(dataset_ids)} datasets "
            f"({', '.join(dataset_ids)}): pooled {built['n_obs']} samples over "
            f"{info['n_genes_corrected']} shared genes and ComBat-corrected the "
            f"matrix ({info['normalization']}) with `dataset` as the batch key. "
            f"Score it with dataset_score_bulk_samples."
        ),
        "mode": "early",
        "plan_reason": plan.get("reason", ""),
        "datasets": list(dataset_ids),
        "batch_key": "batch",
        "per_batch_n": built["per_batch_n"],
        "n_combined_samples": built["n_obs"],
        "n_shared_features": built["n_vars"],
        "n_genes_corrected": info["n_genes_corrected"],
        "n_genes_dropped_constant_within_batch": info["n_genes_dropped_constant_within_batch"],
        "batch_correction": info["method"],
        "normalization": info["normalization"],
        "batch_corrected": True,
        "output_path": str(out_path.resolve()),
        "expression_path": str(out_path.resolve()),
        "next_step": (
            "Pass output_path as expression_path to dataset_score_bulk_samples "
            "(resource=progeny|collectri|hallmark) for per-sample activity scores; the "
            "scores are batch-corrected, so cohort differences no longer confound them."
        ),
        "interpretation_note": (
            "Pooled + ComBat-corrected for per-sample activity scoring: the cohorts were "
            "concatenated on shared genes and `dataset` batch effects were removed with "
            "ComBat BEFORE scoring (no design-matrix covariate exists for per-sample "
            "scoring). Report the datasets pooled, the per-dataset sample counts, and that "
            "ComBat was applied. CAVEAT: ComBat ran without a biological covariate, so if a "
            "group of interest is confounded with cohort, real between-cohort biology may "
            "also be removed. Do NOT use this matrix for differential-expression testing — "
            "for an early DE contrast use decoupler_integrate_datasets (batch covariate)."
        ),
        "artifacts": [
            {
                "description": "Pooled, ComBat batch-corrected AnnData (h5ad) — score this",
                "path": str(out_path.resolve()),
            }
        ],
    }