edgepython_server.py

# This code was written by Claude (Anthropic). The project was directed by Lior Pachter.
"""
edgePython MCP Server — Differential expression analysis via tool calls.

Wraps edgePython's core pipeline so AI agents can run RNA-seq DE analysis
step-by-step: load data → filter → normalize → estimate dispersions →
fit model → test contrasts → get results → generate plots.

Usage:
    python3 edgepython_server.py
"""

import os
import sys
import re
import warnings
import numpy as np
import pandas as pd
from typing import Optional

from fastmcp import FastMCP
from statsmodels.stats.multitest import multipletests

# Ensure edgePython is importable
_edgepython_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if _edgepython_root not in sys.path:
    sys.path.insert(0, _edgepython_root)

import edgepython as ep

# ---------------------------------------------------------------------------
# Server + state
# ---------------------------------------------------------------------------

mcp = FastMCP(
    "edgepython",
    instructions=(
        "RNA-seq differential expression analysis server powered by edgePython. "
        "Bulk workflow: load_data → filter_genes → normalize → set_design → "
        "estimate_dispersion (or estimate_glm_robust_dispersion) → "
        "voom_transform (optional) → fit_model → test_contrast → "
        "get_top_genes / generate_plot / gene_set_test. Quantifier output "
        "(salmon, kallisto, lr-kallisto, oarfish, RSEM): load_quant. kb count "
        "(kallisto | bustools) runs, one per sample: load_kb, as single cells or "
        "pseudo-bulk. DTU: "
        "load_quant → estimate_dispersion → dtu_diff_splice_dge or "
        "dtu_splice_variants. Gene-level p-values from a transcript-level test: "
        "aggregate_to_genes (Lancaster). Methylation: load_bismark → fit_bin_model → "
        "test_bin → get_top_genes. "
        "Single-cell workflow: load_sc_data (or load_kb) → set_sc_design → fit_sc_model → "
        "test_sc_coef → get_sc_top_genes. "
        "Use describe() or describe_sc() to see current pipeline state."
    ),
)

# Module-level state — mutated by tools across a session.
_state: dict = {
    "dgelist": None,       # Current DGEList
    "design": None,        # Design matrix (ndarray)
    "design_info": None,   # Column names for the design matrix
    "fit": None,           # Fitted QL-GLM (DGEGLM-like dict)
    "glm_fit": None,       # Fitted GLM (non-QL)
    "results": {},         # name → DGELRT-like dict
    "last_result": None,   # Name of most recent test result
    "voom": None,          # Most recent voom/voomLmFit output dict
    "filtered": False,
    "normalized": False,
    "dispersions_estimated": False,
    "feature_level": "unknown",  # inferred: gene | transcript | unknown
    "feature_level_note": None,
    "analysis_context": "unknown",  # inferred: bulk | single_cell_10x | single_cell | unknown
    "analysis_context_note": None,
}


def _require(key: str, label: str):
    """Raise if a pipeline stage hasn't been completed yet."""
    if _state[key] is None:
        raise ValueError(
            f"No {label} available. Run the appropriate step first."
        )


def _gene_name(row):
    """Extract gene name from a top_tags result row."""
    # top_tags prepends gene annotation columns; check for GeneID first
    for col in ("GeneID", "gene_id", "gene_name", "Symbol"):
        if col in row.index:
            return str(row[col])
    # Fall back to row index
    return str(row.name)


def _infer_feature_level(feature_ids: Optional[list]) -> tuple[str, str]:
    """Heuristically infer feature level from row IDs."""
    if not feature_ids:
        return "unknown", "unknown (no feature IDs found)"

    ids = [str(x).strip() for x in feature_ids if str(x).strip()]
    if not ids:
        return "unknown", "unknown (no non-empty feature IDs found)"

    # Sample up to 500 IDs for a lightweight heuristic.
    ids = ids[:500]
    n = len(ids)

    tx_patterns = (
        r"(^|[|_.-])ENST\d+",
        r"(^|[|_.-])NM_\d+",
        r"(^|[|_.-])NR_\d+",
        r"(^|[|_.-])XM_\d+",
        r"(^|[|_.-])XR_\d+",
        r"\btx\d+\b",
        r"\btranscript\b",
        r"\bisoform\b",
    )
    gene_patterns = (
        r"(^|[|_.-])ENSG\d+",
        r"(^|[|_.-])FBgn\d+",
        r"(^|[|_.-])WBGene\d+",
    )

    tx_hits = 0
    gene_hits = 0
    for s in ids:
        if any(re.search(p, s, flags=re.IGNORECASE) for p in tx_patterns):
            tx_hits += 1
        if any(re.search(p, s, flags=re.IGNORECASE) for p in gene_patterns):
            gene_hits += 1

    tx_frac = tx_hits / n
    gene_frac = gene_hits / n

    # Bias to gene-level when uncertain (edgePython default use case).
    if tx_frac >= 0.10 and tx_hits >= 5:
        return "transcript", f"transcript/isoform-like IDs ({tx_hits}/{n})"
    if gene_frac >= 0.10 and gene_hits >= 5:
        return "gene", f"gene-like IDs ({gene_hits}/{n})"
    return "gene", "gene-level assumed (IDs do not strongly indicate transcript-level)"


def _infer_analysis_context(
    source: Optional[str] = None,
    data_input: Optional[object] = None,
    sample_names: Optional[list] = None,
    default_bulk: bool = False,
) -> tuple[str, str]:
    """Heuristically infer sequencing/analysis context."""
    src = (source or "").strip().lower()

    # Flatten potential path-like clues.
    clues = []
    if isinstance(data_input, (list, tuple)):
        clues.extend([str(x).lower() for x in data_input])
    elif data_input is not None:
        clues.append(str(data_input).lower())

    clue_text = " ".join(clues)

    # kb count output: cells, or runs summed over cells
    if src == "kb":
        return "single_cell", "kb count output"
    if src == "kb_pseudobulk":
        return "pseudobulk", "kb count runs summed over cells"

    # Strong single-cell/10x signals.
    if src == "10x" or "10x" in clue_text or "cellranger" in clue_text:
        return "single_cell_10x", "10x/cellranger-like source/path"

    if src == "anndata" or any(s.endswith('.h5ad') for s in clues):
        # Could be 10x-derived or non-10x single-cell.
        if "10x" in clue_text or "cellranger" in clue_text:
            return "single_cell_10x", "AnnData with 10x/cellranger-like clues"
        return "single_cell", "AnnData source/path"

    # SMART-seq / random-primed style clues -> bulk-like transcriptome input.
    if any(k in clue_text for k in ("smartseq", "smart-seq", "smart_seq", "random primed", "random_primed", "random-primed")):
        return "bulk", "SMART-seq/random-primed-like clues"

    # Source families typically used for bulk transcript quantification.
    if src in {"kallisto", "salmon", "rsem", "oarfish", "table", "matrix", "dataframe", "sparse", "rds"}:
        return "bulk", f"source='{src}' typically bulk/pseudobulk-like"

    # Sample-name heuristic for 10x-like barcodes.
    if sample_names:
        names = [str(x) for x in sample_names]
        if names:
            barcode_hits = 0
            for n in names:
                s = n.strip()
                if re.fullmatch(r"[ACGT]{8,}(?:-\d+)?", s):
                    barcode_hits += 1
            frac = barcode_hits / len(names)
            if len(names) >= 100 and frac >= 0.5:
                return "single_cell_10x", f"sample names look like 10x barcodes ({barcode_hits}/{len(names)})"

    if default_bulk:
        return "bulk", "bulk assumed from count-matrix workflow"

    return "unknown", "insufficient clues"


def _is_bulk_like_context() -> bool:
    return _state.get("analysis_context") == "bulk"


def _gene_level_warning() -> Optional[str]:
    """Soft warning shown only for likely bulk/random-primed contexts."""
    if _state.get("feature_level") == "gene" and _is_bulk_like_context():
        return (
            "Note: This appears to be gene-level DE in a bulk/random-primed-like context. "
            "If transcript/isoform quantifications are available, consider isoform-level "
            "analysis to avoid masking isoform-specific changes."
        )
    return None


def _append_gene_level_warning(lines: list) -> None:
    warn = _gene_level_warning()
    if warn:
        lines.extend(["", warn])


def _resolve_gene_exon_ids(
    d: dict,
    gene_column: Optional[str] = None,
    exon_column: Optional[str] = None,
):
    """Resolve gene/exon IDs from DGEList gene annotation or fallback IDs."""
    genes = d.get("genes")
    if genes is None:
        gene_ids = d.get("_gene_ids")
        if gene_ids is None:
            raise ValueError("No gene annotation available; provide DGEList genes with GeneID.")
        gene_ids = np.asarray(gene_ids)
        exon_ids = np.arange(len(gene_ids))
        return gene_ids, exon_ids

    gdf = genes if isinstance(genes, pd.DataFrame) else pd.DataFrame(genes)

    if gene_column is not None:
        if gene_column not in gdf.columns:
            raise ValueError(f"gene_column '{gene_column}' not found in gene annotation.")
        gene_ids = gdf[gene_column].astype(str).values
    else:
        found = None
        for col in ("GeneID", "geneid", "gene_id", "Gene", "EnsG"):
            if col in gdf.columns:
                found = col
                break
        if found is None:
            raise ValueError("Could not infer gene IDs. Set gene_column explicitly.")
        gene_ids = gdf[found].astype(str).values

    if exon_column is not None:
        if exon_column not in gdf.columns:
            raise ValueError(f"exon_column '{exon_column}' not found in gene annotation.")
        exon_ids = gdf[exon_column].astype(str).values
    else:
        found = None
        for col in ("ExonID", "exonid", "exon_id", "TranscriptID", "transcript_id"):
            if col in gdf.columns:
                found = col
                break
        if found is not None:
            exon_ids = gdf[found].astype(str).values
        elif not isinstance(gdf.index, pd.RangeIndex):
            exon_ids = gdf.index.astype(str).values
        else:
            exon_ids = np.arange(len(gene_ids))

    return np.asarray(gene_ids), np.asarray(exon_ids)


# ---------------------------------------------------------------------------
# Tool 1: load_data
# ---------------------------------------------------------------------------

@mcp.tool()
def load_data(
    counts_path: str,
    sample_info_path: Optional[str] = None,
    group_column: Optional[str] = None,
    separator: Optional[str] = None,
) -> str:
    """Load a count matrix (and optional sample metadata) to create a DGEList.

    Args:
        counts_path: Path to count matrix file (TSV/CSV). Genes as rows,
            samples as columns. First column is gene IDs.
        sample_info_path: Optional path to sample metadata file (TSV/CSV).
            Must have a column matching sample names in the count matrix.
        group_column: Column in sample info to use as group factor.
            Auto-detected if not provided.
        separator: Column separator ('\\t' or ','). Auto-detected from
            file extension if not provided.
    """
    # Auto-detect separator
    if separator is None:
        separator = "," if counts_path.endswith(".csv") else "\t"

    # Read counts
    counts_df = pd.read_csv(counts_path, sep=separator, index_col=0)
    gene_ids = list(counts_df.index)
    sample_names = list(counts_df.columns)
    counts = counts_df.values.astype(np.float64)

    genes_df = pd.DataFrame({"GeneID": gene_ids})

    # Read sample info if provided
    group = None
    samples_df = None
    if sample_info_path is not None:
        sep2 = "," if sample_info_path.endswith(".csv") else "\t"
        samples_df = pd.read_csv(sample_info_path, sep=sep2)

        # Try to align rows to sample order
        for col in samples_df.columns:
            vals = [str(v) for v in samples_df[col]]
            if set(sample_names).issubset(set(vals)):
                samples_df = samples_df.set_index(col).loc[sample_names].reset_index()
                break

        # Extract group
        if group_column is not None:
            group = list(samples_df[group_column].astype(str))
        else:
            # Auto-detect: first column with few unique values
            for col in samples_df.columns:
                n_unique = samples_df[col].nunique()
                if 2 <= n_unique <= len(sample_names) // 2:
                    group = list(samples_df[col].astype(str))
                    group_column = col
                    break

    dgelist = ep.make_dgelist(
        counts=counts,
        group=group,
        genes=genes_df,
        samples=samples_df,
    )
    # Attach sample/gene names for display
    dgelist["_sample_names"] = sample_names
    dgelist["_gene_ids"] = gene_ids

    level, level_note = _infer_feature_level(gene_ids)
    context, context_note = _infer_analysis_context(
        source=None,
        data_input=counts_path,
        sample_names=sample_names,
        default_bulk=True,
    )

    _state["dgelist"] = dgelist
    _state["feature_level"] = level
    _state["feature_level_note"] = level_note
    _state["analysis_context"] = context
    _state["analysis_context_note"] = context_note
    _state["filtered"] = False
    _state["normalized"] = False
    _state["dispersions_estimated"] = False
    _state["design"] = None
    _state["design_info"] = None
    _state["fit"] = None
    _state["glm_fit"] = None
    _state["results"] = {}
    _state["last_result"] = None
    _state["voom"] = None

    # Build summary
    lib_sizes = dgelist["samples"]["lib.size"].values
    lines = [
        f"Loaded {counts.shape[0]:,} genes × {counts.shape[1]} samples",
        f"Samples: {', '.join(sample_names)}",
    ]
    if group is not None:
        from collections import Counter
        gc = Counter(group)
        lines.append(f"Groups ({group_column}): " + ", ".join(
            f"{k} (n={v})" for k, v in gc.items()
        ))
    lines.append(
        f"Library sizes: {int(lib_sizes.min()):,} – {int(lib_sizes.max()):,} "
        f"(median {int(np.median(lib_sizes)):,})"
    )
    if _state.get("feature_level_note"):
        lines.append(f"Feature level: {_state['feature_level']} ({_state['feature_level_note']})")
    if _state.get("analysis_context_note"):
        lines.append(f"Context: {_state['analysis_context']} ({_state['analysis_context_note']})")
    _append_gene_level_warning(lines)
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 1b: load_data_auto
# ---------------------------------------------------------------------------

@mcp.tool()
def load_data_auto(
    data: Optional[str] = None,
    data_list: Optional[list] = None,
    source: Optional[str] = None,
    format: Optional[str] = None,
    path: Optional[str] = None,
    group: Optional[list] = None,
    labels: Optional[list] = None,
    columns: Optional[list] = None,
    sep: str = "\t",
    obs_col: Optional[str] = None,
    layer: Optional[str] = None,
    ngibbs: int = 100,
    divide: bool = False,
    verbose: bool = True,
) -> str:
    """Load data using edgePython's flexible read_data() loader.

    Args:
        data: Path to file/dir or other supported value. If data_list is provided,
            this is ignored.
        data_list: Optional list of paths (e.g. quantification dirs).
        source: Data source (kallisto, lr_kallisto, salmon, oarfish, rsem,
            10x, kb, table, anndata, rds, sparse, matrix, dataframe).
            Auto-detected if None. For quantifier output with edgeR 5.0's
            annotation and gene-level summaries, use load_quant(); for kb
            count output from several samples, use load_kb().
        format: Source-specific format (e.g. 'h5' or 'tsv' for kallisto).
        path: Base path prefix for relative paths.
        group: Optional group assignments.
        labels: Optional sample labels.
        columns: For table format: (gene_id_col, count_col) 0-based indices.
        sep: Field separator for table/CSV files.
        obs_col: For AnnData: column in .obs to use as group factor.
        layer: For AnnData: layer name to use instead of .X. For kb: the
            nac workflow matrix ('mature', 'nascent', 'ambiguous', ...).
        ngibbs: For RSEM: number of Gibbs samples per sample.
        divide: For quantifier output: divide counts by the estimated
            overdispersions. Default False.
        verbose: Print progress messages.
    """
    if data_list is not None:
        data_input = data_list
    elif data is not None:
        data_input = data
    else:
        raise ValueError("Provide data or data_list.")

    dgelist = ep.read_data(
        data_input,
        source=source,
        format=format,
        path=path,
        group=group,
        labels=labels,
        columns=columns,
        sep=sep,
        obs_col=obs_col,
        layer=layer,
        ngibbs=ngibbs,
        divide=divide,
        verbose=verbose,
    )

    if source is None and isinstance(data_input, str) and os.path.isdir(data_input) \
            and ep.io._is_kb_dir(os.listdir(data_input)):
        source = "kb"
    lines = _register_loaded(dgelist, source=source, data_input=data_input)
    return "\n".join(lines)


def _feature_ids(dgelist) -> Optional[list]:
    """Feature IDs from the gene annotation columns, or its row names."""
    gdf = dgelist.get("genes")
    if gdf is None:
        return None
    for col in ("GeneID", "gene_id", "gene_name", "Symbol"):
        if col in gdf.columns:
            return list(gdf[col].astype(str))
    if not isinstance(gdf.index, pd.RangeIndex):
        return [str(x) for x in gdf.index]
    return None


def _register_loaded(dgelist, source=None, data_input=None) -> list:
    """Make a loaded DGEList the current data, clearing downstream state,
    and return summary lines."""
    _state["dgelist"] = dgelist
    _state["filtered"] = False
    _state["normalized"] = False
    _state["dispersions_estimated"] = False
    _state["design"] = None
    _state["design_info"] = None
    _state["fit"] = None
    _state["glm_fit"] = None
    _state["results"] = {}
    _state["last_result"] = None
    _state["voom"] = None

    counts = dgelist["counts"]
    gene_ids = _feature_ids(dgelist)
    if gene_ids is not None:
        dgelist["_gene_ids"] = gene_ids

    sample_names = None
    if dgelist.get("samples") is not None:
        sdf = dgelist["samples"]
        if sdf.index is not None and len(sdf.index) == counts.shape[1]:
            sample_names = [str(x) for x in sdf.index]
    if sample_names is not None:
        dgelist["_sample_names"] = sample_names

    level, level_note = _infer_feature_level(gene_ids)
    context, context_note = _infer_analysis_context(
        source=source,
        data_input=data_input,
        sample_names=sample_names,
        default_bulk=False,
    )
    _state["feature_level"] = level
    _state["feature_level_note"] = level_note
    _state["analysis_context"] = context
    _state["analysis_context_note"] = context_note

    lines = [f"Loaded {counts.shape[0]:,} genes × {counts.shape[1]} samples"]
    if sample_names is not None:
        lines.append(f"Samples: {', '.join(sample_names)}")
    if _state.get("feature_level_note"):
        lines.append(f"Feature level: {_state['feature_level']} ({_state['feature_level_note']})")
    if _state.get("analysis_context_note"):
        lines.append(f"Context: {_state['analysis_context']} ({_state['analysis_context_note']})")
    _append_gene_level_warning(lines)
    return lines


# ---------------------------------------------------------------------------
# Tool 1c: load_quant
# ---------------------------------------------------------------------------

_QUANT_READERS = {
    "salmon": (ep.catch_salmon, ep.catch_salmon_gene),
    "kallisto": (ep.catch_kallisto, ep.catch_kallisto_gene),
    "lr_kallisto": (ep.catch_lr_kallisto, ep.catch_lr_kallisto_gene),
    "oarfish": (ep.catch_oarfish, ep.catch_oarfish_gene),
    "rsem": (ep.catch_rsem, ep.catch_rsem_gene),
}


@mcp.tool()
def load_quant(
    source: str,
    paths: Optional[list] = None,
    parent_dir: Optional[str] = None,
    level: str = "transcript",
    tx2gene_path: Optional[str] = None,
    divide: bool = True,
    group: Optional[list] = None,
    verbose: bool = False,
) -> str:
    """Load quantifier output with edgeR 5.0's catch* readers.

    Transcript-level counts are divided by the overdispersions estimated from
    bootstrap or Gibbs samples (Baldoni et al. 2024), and Gencode transcript
    names are unpacked into annotation columns (EnsG, TxName, GeneName, ...),
    which dtu_diff_splice_dge and dtu_splice_variants use as gene IDs.
    Gene-level loading sums the counts within genes.

    Args:
        source: 'salmon', 'kallisto', 'lr_kallisto', 'oarfish' or 'rsem'.
        paths: Sample directories (salmon, kallisto), file prefixes (oarfish)
            or isoforms.results files (rsem). Not used for lr_kallisto.
        parent_dir: Directory holding the samples (alternative to paths;
            required for lr_kallisto, holding the quant-tcc output).
        level: 'transcript' (default) or 'gene'.
        tx2gene_path: For level='gene': TSV/CSV whose first two columns map
            transcript IDs to gene IDs, such as kb ref's t2g.txt (headerless;
            its third column gives the gene names). Not needed for Gencode
            transcript names, or for rsem, which uses its genes.results files.
        divide: Divide transcript counts by their overdispersions. Default True.
        group: Optional group assignment for each sample.
        verbose: Print progress messages.
    """
    if source not in _QUANT_READERS:
        raise ValueError(f"source must be one of: {', '.join(_QUANT_READERS)}")
    if level not in ("transcript", "gene"):
        raise ValueError("level must be 'transcript' or 'gene'.")
    tx_reader, gene_reader = _QUANT_READERS[source]
    if source == "lr_kallisto":
        if parent_dir is None:
            raise ValueError("lr_kallisto requires parent_dir.")
        kwargs = {"parent_dir": parent_dir}
    else:
        if paths is None and parent_dir is None:
            raise ValueError("Provide paths or parent_dir.")
        key = {"oarfish": "prefixes", "rsem": "files"}.get(source, "sample_dirs")
        kwargs = {"parent_dir": parent_dir, key: paths}

    if level == "gene":
        if tx2gene_path is not None:
            if source == "rsem":
                raise ValueError("rsem gene-level loading uses its genes.results files; omit tx2gene_path.")
            # A TSV path goes to the reader as it is, read without a header (a
            # header row only adds an unused mapping) so kb's t2g.txt works
            kwargs["tx2gene"] = (pd.read_csv(tx2gene_path) if tx2gene_path.endswith(".csv")
                                 else tx2gene_path)
    with warnings.catch_warnings(record=True) as caught:
        warnings.simplefilter("always")
        if level == "gene":
            dgelist = gene_reader(divide=divide, verbose=verbose, **kwargs)
        else:
            dgelist = tx_reader(divide=divide, verbose=verbose, **kwargs)
    for w in caught:
        if not str(w.message).startswith("Samples were quantified with different"):
            warnings.warn_explicit(w.message, w.category, w.filename, w.lineno)

    if group is not None:
        if len(group) != dgelist["counts"].shape[1]:
            raise ValueError(f"group has {len(group)} values for {dgelist['counts'].shape[1]} samples.")
        dgelist["samples"]["group"] = pd.Categorical([str(g) for g in group])

    lines = _register_loaded(dgelist, source=source, data_input=paths or parent_dir)
    od = dgelist.get("genes")
    if od is not None and "Overdispersion" in od.columns:
        vals = od["Overdispersion"].to_numpy(dtype=float)
        lines.insert(1, (
            f"Overdispersion: median {np.nanmedian(vals):.3g}, "
            f"range {np.nanmin(vals):.3g}–{np.nanmax(vals):.3g}"
            + (" (counts divided)" if dgelist.get("divided.counts") else "")
        ))
    info = dgelist.get("run.info")
    run_lines = []
    for name, r in (info.iterrows() if info is not None else []):
        parts = [f"{r[k]:,.0f} reads" if k == "n_processed" else f"{r[k]}% pseudoaligned"
                 for k in ("n_processed", "p_pseudoaligned") if k in r and pd.notna(r[k])]
        if "kallisto_version" in r and pd.notna(r["kallisto_version"]):
            parts.append(f"kallisto {r['kallisto_version']}")
        if parts:
            run_lines.append(f"  {name}: {', '.join(parts)}")
    if run_lines:
        lines += ["kallisto runs:"] + run_lines
    for w in caught:
        if str(w.message).startswith("Samples were quantified with different"):
            lines.append(f"Warning: {w.message}")
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 2: describe
# ---------------------------------------------------------------------------

@mcp.tool()
def describe() -> str:
    """Report the current state of the analysis pipeline."""
    lines = []

    d = _state["dgelist"]
    if d is None:
        return "No data loaded. Use load_data() first."

    counts = d["counts"]
    lines.append(f"Data: {counts.shape[0]:,} genes × {counts.shape[1]} samples")
    lines.append(f"Samples: {', '.join(d.get('_sample_names', [str(i) for i in range(counts.shape[1])]))}")
    lines.append(f"Filtered: {'yes' if _state['filtered'] else 'no'}")
    lines.append(f"Normalized: {'yes' if _state['normalized'] else 'no'}")

    if _state["normalized"]:
        nf = d["samples"]["norm.factors"].values
        lines.append(f"  Norm factors: {', '.join(f'{v:.4f}' for v in nf)}")

    lines.append(f"Dispersions estimated: {'yes' if _state['dispersions_estimated'] else 'no'}")
    fl = _state.get("feature_level", "unknown")
    note = _state.get("feature_level_note")
    if note:
        lines.append(f"Feature level: {fl} ({note})")
    else:
        lines.append(f"Feature level: {fl}")

    cx = _state.get("analysis_context", "unknown")
    cx_note = _state.get("analysis_context_note")
    if cx_note:
        lines.append(f"Context: {cx} ({cx_note})")
    else:
        lines.append(f"Context: {cx}")
    if _state["dispersions_estimated"]:
        cd = d.get("common.dispersion")
        if cd is not None:
            lines.append(f"  Common dispersion: {cd:.4f} (BCV: {np.sqrt(cd):.4f})")

    if _state["design"] is not None:
        dm = _state["design"]
        lines.append(f"Design: {dm.shape[1]} coefficients — {', '.join(_state['design_info'])}")

    if _state["fit"] is not None:
        fit = _state["fit"]
        pdf = fit.get("df.prior")
        if pdf is not None:
            if np.isscalar(pdf):
                lines.append(f"Model fitted (QL): prior df = {pdf:.2f}")
            else:
                lines.append(f"Model fitted (QL): median prior df = {np.median(pdf):.2f}")
    if _state["voom"] is not None:
        vw = _state["voom"].get("weights")
        if vw is not None:
            lines.append(
                f"Voom weights: {vw.shape[0]:,}×{vw.shape[1]}, "
                f"range {np.nanmin(vw):.3g}–{np.nanmax(vw):.3g}"
            )

    if _state["results"]:
        lines.append(f"Test results: {', '.join(_state['results'].keys())}")
    else:
        lines.append("No test results yet.")

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 2b: reset_state
# ---------------------------------------------------------------------------

@mcp.tool()
def reset_state() -> str:
    """Reset all pipeline state."""
    _state["dgelist"] = None
    _state["design"] = None
    _state["design_info"] = None
    _state["fit"] = None
    _state["glm_fit"] = None
    _state["results"] = {}
    _state["last_result"] = None
    _state["voom"] = None
    _state["filtered"] = False
    _state["normalized"] = False
    _state["dispersions_estimated"] = False
    _state["feature_level"] = "unknown"
    _state["feature_level_note"] = None
    _state["analysis_context"] = "unknown"
    _state["analysis_context_note"] = None
    return "State cleared."


# ---------------------------------------------------------------------------
# Tool 3: filter_genes
# ---------------------------------------------------------------------------

@mcp.tool()
def filter_genes(
    min_count: Optional[float] = None,
    min_total_count: Optional[float] = None,
    preset: Optional[str] = None,
    max_n_samples: Optional[float] = None,
) -> str:
    """Filter out lowly-expressed genes using edgeR's filterByExpr logic.

    Uses adaptive CPM thresholds based on library sizes and group structure.

    Args:
        min_count: Minimum count threshold (used to derive CPM cutoff).
            Default: 10, or as set by preset.
        min_total_count: Minimum total count across all samples. Default: 15,
            or as set by preset.
        preset: edgeR 5.0 preset thresholds: 'stringent', 'moderate' or
            'lenient'. Overrides min_count and min_total_count.
        max_n_samples: Cap on the number of samples a feature must be
            expressed in (edgeR 5.0 default 40).
    """
    _require("dgelist", "DGEList")
    d = _state["dgelist"]

    group = d["samples"].get("group")
    keep = ep.filter_by_expr(
        d, group=group, min_count=min_count, min_total_count=min_total_count,
        preset=preset, max_n_samples=max_n_samples,
    )

    n_before = d["counts"].shape[0]
    n_after = int(keep.sum())

    # Apply filter
    d["counts"] = d["counts"][keep]
    if "genes" in d and d["genes"] is not None:
        d["genes"] = d["genes"].iloc[keep.nonzero()[0]].reset_index(drop=True)
    gene_ids = d.get("_gene_ids")
    if gene_ids is not None:
        d["_gene_ids"] = [gene_ids[i] for i in range(n_before) if keep[i]]

    # Clear downstream state
    _state["dispersions_estimated"] = False
    _state["fit"] = None
    _state["glm_fit"] = None
    _state["results"] = {}
    _state["last_result"] = None
    _state["voom"] = None
    _state["filtered"] = True

    # Recompute lib sizes after filtering
    d["samples"]["lib.size"] = d["counts"].sum(axis=0).astype(np.float64)

    return (
        f"Filtered: {n_before:,} → {n_after:,} genes "
        f"({n_before - n_after:,} removed)"
    )


# ---------------------------------------------------------------------------
# Tool 4: normalize
# ---------------------------------------------------------------------------

@mcp.tool()
def normalize(method: str = "TMM") -> str:
    """Apply library-size normalization.

    Args:
        method: Normalization method. One of 'TMM' (default), 'TMMwsp',
            'RLE', 'upperquartile', or 'none'.
    """
    _require("dgelist", "DGEList")
    d = _state["dgelist"]

    d = ep.calc_norm_factors(d, method=method)
    _state["dgelist"] = d
    _state["normalized"] = True
    _state["voom"] = None

    nf = d["samples"]["norm.factors"].values
    eff = d["samples"]["lib.size"].values * nf

    lines = [
        f"Normalization: {method}",
        "Sample | Norm Factor | Effective Lib Size",
        "-------|------------|-------------------",
    ]
    names = d.get("_sample_names", [f"S{i+1}" for i in range(len(nf))])
    for i, name in enumerate(names):
        lines.append(f"{name} | {nf[i]:.4f} | {int(eff[i]):,}")

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 4b: normalize_chip — ChIP-Seq normalization to input
# ---------------------------------------------------------------------------

@mcp.tool()
def normalize_chip(
    input_csv: str,
    dispersion: float = 0.01,
    niter: int = 6,
    loss: str = "p",
) -> str:
    """Normalize ChIP-Seq counts to input control and test for enrichment.

    Reads input control counts from a CSV file and normalizes the currently
    loaded ChIP response counts against them, setting offsets for downstream
    GLM fitting.

    Args:
        input_csv: Path to a CSV/TSV file of input control counts with the
            same genes (rows) and samples (columns) as the loaded data.
        dispersion: Negative binomial dispersion (default 0.01).
        niter: Number of iterations for scaling-factor estimation.
        loss: Loss function — 'p' (cumulative probabilities) or 'z' (z-value).
    """
    _require("dgelist", "DGEList")
    d = _state["dgelist"]

    sep = "\t" if input_csv.endswith(".tsv") else ","
    inp_df = pd.read_csv(input_csv, sep=sep, index_col=0)
    inp_mat = inp_df.values.astype(np.float64)

    if inp_mat.shape[0] != d["counts"].shape[0]:
        raise ValueError(
            f"Row count mismatch: input has {inp_mat.shape[0]} rows but "
            f"loaded data has {d['counts'].shape[0]} genes."
        )

    d = ep.calc_norm_offsets_for_chip(
        inp_mat, d, dispersion=dispersion, niter=niter, loss=loss,
    )
    _state["dgelist"] = d
    _state["voom"] = None

    lines = [
        f"ChIP normalization applied (dispersion={dispersion}, loss='{loss}').",
        f"Offset matrix set ({d['offset'].shape[0]:,} genes × "
        f"{d['offset'].shape[1]} samples).",
    ]
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 4c: chip_enrichment_test — per-sample ChIP enrichment p-values
# ---------------------------------------------------------------------------

@mcp.tool()
def chip_enrichment_test(
    input_csv: str,
    sample: int = 0,
    dispersion: float = 0.01,
    niter: int = 6,
    n_top: int = 20,
) -> str:
    """Test for ChIP-Seq enrichment relative to input for one sample.

    Returns the top enriched features ranked by p-value.

    Args:
        input_csv: Path to input control counts (CSV/TSV).
        sample: Sample index (0-based) to test.
        dispersion: Negative binomial dispersion.
        niter: Number of iterations.
        n_top: Number of top enriched features to report.
    """
    _require("dgelist", "DGEList")
    d = _state["dgelist"]

    sep = "\t" if input_csv.endswith(".tsv") else ","
    inp_df = pd.read_csv(input_csv, sep=sep, index_col=0)
    inp_mat = inp_df.values.astype(np.float64)

    if sample < 0 or sample >= d["counts"].shape[1]:
        raise ValueError(f"sample must be 0..{d['counts'].shape[1]-1}")

    out = ep.normalize_chip_to_input(
        inp_mat[:, sample], d["counts"][:, sample],
        dispersion=dispersion, niter=niter,
    )

    gene_ids = d.get("_gene_ids")
    if gene_ids is None:
        gene_ids = [str(i) for i in range(d["counts"].shape[0])]

    order = np.argsort(out['p_value'])
    n_show = min(n_top, len(order))

    names = d.get("_sample_names", [f"S{i}" for i in range(d["counts"].shape[1])])
    lines = [
        f"ChIP enrichment test — sample '{names[sample]}'",
        f"Scaling factor: {out['scaling_factor']:.4f}",
        f"Proportion enriched: {out['prop_enriched']:.4f}",
        "",
        f"Top {n_show} enriched features:",
        "Gene | p-value | mid-p",
        "-----|---------|------",
    ]
    for i in range(n_show):
        idx = order[i]
        lines.append(
            f"{gene_ids[idx]} | {out['p_value'][idx]:.4e} | "
            f"{out['pmid_value'][idx]:.4e}"
        )

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 5: set_design
# ---------------------------------------------------------------------------

@mcp.tool()
def set_design(formula: str) -> str:
    """Set the experimental design using an R-style formula.

    Args:
        formula: R-style formula, e.g. '~ 0 + group', '~ group + batch'.
            Variables must match column names in the sample metadata.
    """
    _require("dgelist", "DGEList")
    d = _state["dgelist"]

    # Build data for patsy
    samples = d.get("samples")
    if samples is None:
        raise ValueError("No sample metadata — load sample info with load_data().")

    design = ep.model_matrix(formula, data=samples)

    # Extract column names from patsy
    import patsy
    dinfo = patsy.dmatrix(formula, data=samples, return_type="dataframe")
    col_names = list(dinfo.columns)

    _state["design"] = design
    _state["design_info"] = col_names
    # Clear downstream
    _state["fit"] = None
    _state["glm_fit"] = None
    _state["results"] = {}
    _state["last_result"] = None
    _state["voom"] = None

    lines = [
        f"Design matrix: {design.shape[0]} samples × {design.shape[1]} coefficients",
        f"Columns: {', '.join(col_names)}",
        f"Rank: {np.linalg.matrix_rank(design)}",
    ]
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 5b: set_design_matrix
# ---------------------------------------------------------------------------

@mcp.tool()
def set_design_matrix(matrix: list, columns: list) -> str:
    """Set the design matrix directly.

    Args:
        matrix: 2D list of floats with shape (samples, coefficients).
        columns: Column names for the design matrix.
    """
    _require("dgelist", "DGEList")
    d = _state["dgelist"]
    counts = d["counts"]

    design = np.asarray(matrix, dtype=np.float64)
    if design.ndim != 2:
        raise ValueError("Design matrix must be 2D.")
    if design.shape[0] != counts.shape[1]:
        raise ValueError(
            f"Design has {design.shape[0]} rows but data has {counts.shape[1]} samples."
        )
    if len(columns) != design.shape[1]:
        raise ValueError(
            f"Design has {design.shape[1]} columns but {len(columns)} column names provided."
        )

    _state["design"] = design
    _state["design_info"] = list(columns)
    _state["fit"] = None
    _state["glm_fit"] = None
    _state["results"] = {}
    _state["last_result"] = None
    _state["voom"] = None

    lines = [
        f"Design matrix set: {design.shape[0]} samples × {design.shape[1]} coefficients",
        f"Columns: {', '.join(columns)}",
        f"Rank: {np.linalg.matrix_rank(design)}",
    ]
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 6: estimate_dispersion
# ---------------------------------------------------------------------------

@mcp.tool()
def estimate_dispersion(robust: bool = True) -> str:
    """Estimate NB dispersions (common, trended, tagwise).

    Args:
        robust: Use robust empirical Bayes. Default True.
    """
    _require("dgelist", "DGEList")
    d = _state["dgelist"]

    kwargs = {"robust": robust}
    if _state["design"] is not None:
        kwargs["design"] = _state["design"]

    d = ep.estimate_disp(d, **kwargs)
    _state["dgelist"] = d
    _state["dispersions_estimated"] = True
    _state["voom"] = None

    cd = d.get("common.dispersion", None)
    prior_df = d.get("prior.df", None)
    lines = []
    if cd is not None:
        lines.append(f"Common dispersion: {cd:.4f}")
        lines.append(f"BCV: {np.sqrt(cd):.4f}")
    if prior_df is not None:
        if np.isscalar(prior_df):
            lines.append(f"Prior df: {prior_df:.2f}")
        else:
            lines.append(f"Prior df (median): {np.median(prior_df):.2f}")
    td = d.get("trended.dispersion")
    if td is not None:
        lines.append(f"Trended dispersion range: {td.min():.4f} – {td.max():.4f}")

    return "\n".join(lines) if lines else "Dispersions estimated."


# ---------------------------------------------------------------------------
# Tool 6b: estimate_glm_robust_dispersion
# ---------------------------------------------------------------------------

@mcp.tool()
def estimate_glm_robust_dispersion(
    prior_df: float = 10.0,
    update_trend: bool = True,
    maxit: int = 6,
    k: float = 1.345,
    residual_type: str = "pearson",
    verbose: bool = False,
) -> str:
    """Estimate robust GLM dispersions with iterative Huber reweighting.

    Requires a design matrix; mirrors edgeR's estimateGLMRobustDisp.
    """
    _require("dgelist", "DGEList")
    _require("design", "design matrix")
    d = _state["dgelist"]

    d = ep.estimate_glm_robust_disp(
        d,
        design=_state["design"],
        prior_df=float(prior_df),
        update_trend=bool(update_trend),
        maxit=int(maxit),
        k=float(k),
        residual_type=str(residual_type),
        verbose=bool(verbose),
        record=False,
    )
    _state["dgelist"] = d
    _state["dispersions_estimated"] = True
    _state["voom"] = None

    lines = ["Robust GLM dispersions estimated (Huber reweighting)."]
    cd = d.get("common.dispersion")
    if cd is not None:
        lines.append(f"Common dispersion: {float(cd):.4f} (BCV {np.sqrt(float(cd)):.4f})")
    td = d.get("trended.dispersion")
    if td is not None:
        lines.append(f"Trended dispersion range: {np.min(td):.4f} – {np.max(td):.4f}")
    tw = d.get("tagwise.dispersion")
    if tw is not None:
        lines.append(f"Tagwise dispersion range: {np.min(tw):.4f} – {np.max(tw):.4f}")
    w = d.get("weights")
    if w is not None:
        lines.append(f"Robust weights range: {np.min(w):.3g} – {np.max(w):.3g}")
    _append_gene_level_warning(lines)
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 6c: voom_transform
# ---------------------------------------------------------------------------

@mcp.tool()
def voom_transform(
    span: float = 0.5,
    adaptive_span: bool = True,
    sample_weights: bool = False,
    block_column: Optional[str] = None,
) -> str:
    """Run voom/voomLmFit-style mean-variance weighting.

    Uses current counts + design and stores the voom result in MCP state.

    Args:
        span: Lowess span, used when adaptive_span is False.
        adaptive_span: Choose the span from the number of genes, as limma
            3.99 does by default. Default True.
        sample_weights: Estimate sample quality weights.
        block_column: Sample metadata column giving blocks for
            duplicateCorrelation.
    """
    _require("dgelist", "DGEList")
    _require("design", "design matrix")
    d = _state["dgelist"]

    samples = d.get("samples")
    lib_size = None
    if samples is not None and "lib.size" in samples:
        lib_size = np.asarray(samples["lib.size"].values, dtype=np.float64)
        if "norm.factors" in samples:
            lib_size = lib_size * np.asarray(samples["norm.factors"].values, dtype=np.float64)

    block = None
    if block_column is not None:
        if samples is None or block_column not in samples.columns:
            raise ValueError(f"block_column '{block_column}' not found in sample metadata.")
        block = samples[block_column].astype(str).values

    v = ep.voom_lmfit(
        counts=np.asarray(d["counts"], dtype=np.float64),
        design=np.asarray(_state["design"], dtype=np.float64),
        block=block,
        sample_weights=bool(sample_weights),
        lib_size=lib_size,
        span=float(span),
        adaptive_span=bool(adaptive_span),
        normalize_method="none",
        save_plot=False,
        keep_elist=True,
    )
    _state["voom"] = v

    w = v.get("weights")
    lines = [
        "voom transform completed.",
        f"Input: {d['counts'].shape[0]:,} features × {d['counts'].shape[1]} samples",
        f"Span used: {float(v.get('span', span)):.3f}",
    ]
    if w is not None:
        lines.append(
            f"Weights: shape {w.shape[0]:,}×{w.shape[1]}, "
            f"range {np.nanmin(w):.3g}–{np.nanmax(w):.3g}"
        )
    corr = v.get("correlation")
    if corr is not None and np.isfinite(corr):
        lines.append(f"Consensus block correlation: {float(corr):.4f}")
    sw = v.get("sample_weights")
    if sw is not None:
        lines.append(f"Sample-weight range: {np.nanmin(sw):.3g}–{np.nanmax(sw):.3g}")
    _append_gene_level_warning(lines)
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 7: fit_model
# ---------------------------------------------------------------------------

@mcp.tool()
def fit_model(robust: bool = True, sample_weights: bool = False) -> str:
    """Fit a quasi-likelihood negative binomial GLM.

    Args:
        robust: Use robust empirical Bayes for QL dispersion. Default True.
        sample_weights: Estimate sample quality weights from the fit (edgeR
            5.0's sampleWeights) and refit with them. Default False.
    """
    _require("dgelist", "DGEList")
    _require("design", "design matrix")
    d = _state["dgelist"]

    fit = ep.glm_ql_fit(d, design=_state["design"], robust=robust)
    sw = None
    if sample_weights:
        sw = ep.sample_weights(fit)
        fit = ep.glm_ql_fit(d, design=_state["design"], robust=robust, weights=sw)
    _state["fit"] = fit
    # Clear old results
    _state["results"] = {}
    _state["last_result"] = None
    _state["voom"] = None

    n_genes = fit["coefficients"].shape[0]
    n_coef = fit["coefficients"].shape[1]
    pdf = fit.get("df.prior")
    s2 = fit.get("s2.prior")

    lines = [f"Fitted QL-GLM: {n_genes:,} genes, {n_coef} coefficients"]
    if pdf is not None:
        if np.isscalar(pdf):
            lines.append(f"Prior df: {pdf:.2f}")
        else:
            lines.append(f"Prior df (median): {np.median(pdf):.2f}")
    if s2 is not None:
        if np.isscalar(s2):
            lines.append(f"Prior QL dispersion: {s2:.4f}")
        else:
            lines.append(f"Prior QL dispersion (median): {np.median(s2):.4f}")
    if sw is not None:
        names = d.get("_sample_names") or [f"S{i}" for i in range(sw.shape[1])]
        lines.append("Sample weights: " + ", ".join(
            f"{n} {w:.3f}" for n, w in zip(names, sw[0])
        ))
    lines.append(f"Coefficients: {', '.join(_state['design_info'])}")

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 7b: fit_glm
# ---------------------------------------------------------------------------

@mcp.tool()
def fit_glm() -> str:
    """Fit a negative binomial GLM (non-QL)."""
    _require("dgelist", "DGEList")
    _require("design", "design matrix")
    d = _state["dgelist"]

    fit = ep.glm_fit(d, design=_state["design"])
    _state["glm_fit"] = fit
    _state["results"] = {}
    _state["last_result"] = None
    _state["voom"] = None

    n_genes = fit["coefficients"].shape[0]
    n_coef = fit["coefficients"].shape[1]
    lines = [f"Fitted GLM: {n_genes:,} genes, {n_coef} coefficients"]
    lines.append(f"Coefficients: {', '.join(_state['design_info'])}")
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 8: test_contrast
# ---------------------------------------------------------------------------

def _parse_contrast(contrast_str: str, col_names: list) -> np.ndarray:
    """Parse a contrast string like 'A - B' into a numeric vector.

    Supports:
        'A - B'           → simple difference
        '(A + B)/2 - C'   → average vs reference
        'A'               → single coefficient test
    """
    contrast = np.zeros(len(col_names))

    # Build lookup map: multiple aliases for each column index.
    # Patsy generates names like 'group[A]', 'group[T.A]', 'Intercept'.
    # Users may type 'groupA', 'group[A]', 'A', etc.
    clean_map = {}
    for i, cn in enumerate(col_names):
        clean_map[cn] = i  # exact name, e.g. 'group[A]'

        # Strip patsy treatment-coded wrapper: 'group[T.xxx]' → 'xxx'
        m = re.match(r'^(.+)\[T\.\s*(.+)\]$', cn)
        if m:
            prefix, val = m.group(1), m.group(2)
            clean_map[val] = i              # 'xxx'
            clean_map[prefix + val] = i     # 'groupxxx'

        # Strip patsy no-intercept wrapper: 'group[xxx]' → 'xxx'
        m2 = re.match(r'^(.+)\[(.+)\]$', cn)
        if m2:
            prefix, val = m2.group(1), m2.group(2)
            clean_map[val] = i              # 'xxx'  (e.g. 'A')
            clean_map[prefix + val] = i     # 'groupA'

    def _find_col(term: str) -> int:
        term = term.strip()
        if term in clean_map:
            return clean_map[term]
        # Case-insensitive fallback
        for k, v in clean_map.items():
            if k.lower() == term.lower():
                return v
        raise ValueError(
            f"Term '{term}' not found in design columns: {col_names}\n"
            f"Recognized aliases: {sorted(clean_map.keys())}"
        )

    # Try simple 'A - B' pattern first
    parts = re.split(r'\s*-\s*', contrast_str)
    if len(parts) == 2 and '(' not in contrast_str and '+' not in contrast_str:
        i_pos = _find_col(parts[0])
        i_neg = _find_col(parts[1])
        contrast[i_pos] = 1
        contrast[i_neg] = -1
        return contrast

    # General parser: split into +/- terms
    # Add '+' at start if not starting with '-'
    s = contrast_str.strip()
    if not s.startswith('-'):
        s = '+' + s

    # Tokenize: find all (+/-)(coefficient_or_group)
    tokens = re.findall(r'([+-])\s*(?:(\d*\.?\d*)\s*\*?\s*)?([A-Za-z_][\w.]*)', s)
    for sign, coeff, name in tokens:
        idx = _find_col(name)
        c = float(coeff) if coeff else 1.0
        if sign == '-':
            c = -c
        contrast[idx] += c

    if np.all(contrast == 0):
        raise ValueError(
            f"Could not parse contrast '{contrast_str}'. "
            f"Available columns: {col_names}"
        )

    return contrast


def _run_ql_test(fit, method: str, lfc: float, coef=None, contrast=None):
    """Run a QL F-test, a QL F-test against a fold-change threshold, or TREAT."""
    if method == "treat":
        return ep.glm_treat(fit, coef=coef, contrast=contrast, lfc=lfc)
    if method == "qlf_threshold":
        return ep.glm_ql_ftest(fit, coef=coef, contrast=contrast, lfc=lfc)
    if method == "qlf":
        return ep.glm_ql_ftest(fit, coef=coef, contrast=contrast)
    raise ValueError("method must be 'qlf', 'qlf_threshold' or 'treat'.")


@mcp.tool()
def test_contrast(
    contrast: str,
    name: Optional[str] = None,
    method: str = "qlf",
    lfc: float = 0.585,
) -> str:
    """Test a contrast for differential expression.

    Args:
        contrast: Contrast string, e.g. 'groupB - groupA'. Terms must
            match design matrix column names (shown by set_design/describe).
        name: Label for this result set. Default: auto-generated.
        method: Test method — 'qlf' (quasi-likelihood F-test, default),
            'qlf_threshold' (edgeR 5.0's QL F-test against a log-FC
            threshold, with UPSHOT p-values) or 'treat' (TREAT).
        lfc: Log2-fold-change threshold for 'qlf_threshold' and 'treat'.
            Default: log2(1.5) ≈ 0.585.
    """
    _require("fit", "fitted model")
    fit = _state["fit"]
    col_names = _state["design_info"]

    contrast_vec = _parse_contrast(contrast, col_names)

    res = _run_ql_test(fit, method, lfc, contrast=contrast_vec)

    if name is None:
        name = contrast.replace(" ", "")

    _state["results"][name] = res
    _state["last_result"] = name

    # Summary
    tt = ep.top_tags(res, n=res["table"].shape[0])
    table = tt["table"]
    fdr = table["FDR"].values
    logfc = table["logFC"].values

    n_up = int(((fdr < 0.05) & (logfc > 0)).sum())
    n_down = int(((fdr < 0.05) & (logfc < 0)).sum())
    n_ns = int((fdr >= 0.05).sum())

    lines = [
        f"Test: {method.upper()} — {contrast}",
        f"DE genes (FDR < 0.05): {n_up} up, {n_down} down, {n_ns} NS",
        "",
        "Top 5 genes:",
    ]

    top5 = ep.top_tags(res, n=5)["table"]
    # Format top genes
    lines.append(f"{'Gene':<20} {'logFC':>8} {'logCPM':>8} {'PValue':>10} {'FDR':>10}")
    lines.append("-" * 60)
    for _, row in top5.iterrows():
        gene = _gene_name(row)
        if len(gene) > 19:
            gene = gene[:17] + ".."
        lines.append(
            f"{gene:<20} {row['logFC']:>8.3f} {row['logCPM']:>8.2f} "
            f"{row['PValue']:>10.2e} {row['FDR']:>10.2e}"
        )

    _append_gene_level_warning(lines)
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 8b: test_coef
# ---------------------------------------------------------------------------

@mcp.tool()
def test_coef(
    coef: int,
    name: Optional[str] = None,
    method: str = "qlf",
    lfc: float = 0.585,
) -> str:
    """Test a single coefficient by index.

    Args:
        coef: 0-based coefficient index.
        name: Label for this result set. Default: auto-generated.
        method: 'qlf' (default), 'qlf_threshold' or 'treat'; see test_contrast.
        lfc: Log2-fold-change threshold for 'qlf_threshold' and 'treat'.
    """
    _require("fit", "fitted model")
    fit = _state["fit"]

    res = _run_ql_test(fit, method, lfc, coef=coef)

    if name is None:
        name = f"coef{coef}"

    _state["results"][name] = res
    _state["last_result"] = name

    tt = ep.top_tags(res, n=res["table"].shape[0])
    table = tt["table"]
    fdr = table["FDR"].values
    logfc = table["logFC"].values

    n_up = int(((fdr < 0.05) & (logfc > 0)).sum())
    n_down = int(((fdr < 0.05) & (logfc < 0)).sum())
    n_ns = int((fdr >= 0.05).sum())

    lines = [
        f"Test: {method.upper()} — coef {coef}",
        f"DE genes (FDR < 0.05): {n_up} up, {n_down} down, {n_ns} NS",
    ]
    _append_gene_level_warning(lines)
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 8c: glm_lrt_test
# ---------------------------------------------------------------------------

@mcp.tool()
def glm_lrt_test(
    coef: Optional[int] = None,
    contrast: Optional[list] = None,
    name: Optional[str] = None,
    use_ql_fit: bool = False,
) -> str:
    """Run a likelihood-ratio test (LRT) from a GLM fit.

    Args:
        coef: 0-based coefficient index (default: last).
        contrast: Optional contrast vector (same length as coefficients).
        name: Label for this result set. Default: auto-generated.
        use_ql_fit: If True, uses the QL fit from fit_model().
    """
    fit = _state["fit"] if use_ql_fit else _state["glm_fit"]
    if fit is None:
        raise ValueError("No GLM fit available. Run fit_glm() first.")

    res = ep.glm_lrt(fit, coef=coef, contrast=contrast)

    if name is None:
        if contrast is not None:
            name = "lrt_contrast"
        elif coef is not None:
            name = f"lrt_coef{coef}"
        else:
            name = "lrt"

    _state["results"][name] = res
    _state["last_result"] = name

    tt = ep.top_tags(res, n=res["table"].shape[0])
    table = tt["table"]
    fdr = table["FDR"].values
    logfc = table["logFC"].values if "logFC" in table.columns else np.zeros(len(table))

    n_up = int(((fdr < 0.05) & (logfc > 0)).sum())
    n_down = int(((fdr < 0.05) & (logfc < 0)).sum())
    n_ns = int((fdr >= 0.05).sum())

    lines = [
        "Test: LRT",
        f"DE genes (FDR < 0.05): {n_up} up, {n_down} down, {n_ns} NS",
    ]
    _append_gene_level_warning(lines)
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 8d: exact_test
# ---------------------------------------------------------------------------

@mcp.tool()
def exact_test(
    pair: Optional[list] = None,
    dispersion: str = "auto",
    rejection_region: str = "doubletail",
    big_count: int = 900,
    prior_count: float = 0.125,
    name: Optional[str] = None,
) -> str:
    """Run edgeR-style exact test for two-group DE.

    Args:
        pair: Two group labels to compare. Default: first two groups.
        dispersion: 'auto', 'common', 'trended', 'tagwise', or numeric.
        rejection_region: 'doubletail', 'deviance', or 'smallp'.
        big_count: Threshold for beta approximation.
        prior_count: Prior count for logFC calculation.
        name: Label for this result set. Default: auto-generated.
    """
    _require("dgelist", "DGEList")
    d = _state["dgelist"]

    res = ep.exact_test(
        d,
        pair=pair,
        dispersion=dispersion,
        rejection_region=rejection_region,
        big_count=big_count,
        prior_count=prior_count,
    )

    if name is None:
        if pair is None:
            name = "exact"
        else:
            name = f"exact_{pair[0]}_vs_{pair[1]}"

    _state["results"][name] = res
    _state["last_result"] = name

    tt = ep.top_tags(res, n=res["table"].shape[0])
    table = tt["table"]
    fdr = table["FDR"].values
    logfc = table["logFC"].values

    n_up = int(((fdr < 0.05) & (logfc > 0)).sum())
    n_down = int(((fdr < 0.05) & (logfc < 0)).sum())
    n_ns = int((fdr >= 0.05).sum())

    lines = [
        "Test: Exact",
        f"DE genes (FDR < 0.05): {n_up} up, {n_down} down, {n_ns} NS",
    ]
    _append_gene_level_warning(lines)
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 8e: dtu_diff_splice_dge
# ---------------------------------------------------------------------------

@mcp.tool()
def dtu_diff_splice_dge(
    gene_column: Optional[str] = None,
    exon_column: Optional[str] = None,
    prior_count: float = 0.125,
    name: str = "dtu_diff_splice_dge",
) -> str:
    """Run DTU testing with diff_splice_dge on the current DGEList.

    Requires a two-group design in sample metadata (samples['group']) and
    dispersions from estimate_dispersion(). Gene IDs are taken from the gene
    annotation (GeneID, gene_id or the EnsG column written by load_quant).
    """
    _require("dgelist", "DGEList")
    d = _state["dgelist"]
    if not _state["dispersions_estimated"]:
        raise ValueError("No dispersions available. Run estimate_dispersion() first.")

    if d.get("samples") is None or "group" not in d["samples"].columns:
        raise ValueError("DGEList samples['group'] is required for diff_splice_dge.")

    gene_ids, exon_ids = _resolve_gene_exon_ids(d, gene_column=gene_column, exon_column=exon_column)
    res = ep.diff_splice_dge(
        d,
        geneid=gene_ids,
        exonid=exon_ids,
        group=d["samples"]["group"].values,
        dispersion="auto",
        prior_count=float(prior_count),
    )

    _state["results"][name] = res
    _state["last_result"] = name

    gt = res["gene.table"].copy()
    sig = gt[gt["FDR"] < 0.05]
    top = gt.sort_values("FDR", ascending=True).head(5)
    lines = [
        f"DTU: diff_splice_dge completed ({len(gt):,} genes).",
        f"Genes with DTU signal (FDR < 0.05): {len(sig):,}",
        "",
        "Top 5 genes by FDR:",
        f"{'GeneID':<24} {'NExons':>8} {'PValue':>12} {'FDR':>12}",
        "-" * 62,
    ]
    for _, row in top.iterrows():
        gid = str(row["GeneID"])
        if len(gid) > 23:
            gid = gid[:21] + ".."
        lines.append(
            f"{gid:<24} {int(row['NExons']):>8d} "
            f"{float(row['PValue']):>12.2e} {float(row['FDR']):>12.2e}"
        )
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 8f: dtu_splice_variants
# ---------------------------------------------------------------------------

@mcp.tool()
def dtu_splice_variants(
    gene_column: Optional[str] = None,
    name: str = "dtu_splice_variants",
) -> str:
    """Run splice_variants on the current DGEList to flag genes with isoform heterogeneity."""
    _require("dgelist", "DGEList")
    d = _state["dgelist"]
    gene_ids, _ = _resolve_gene_exon_ids(d, gene_column=gene_column, exon_column=None)

    sv = ep.splice_variants(d, gene_ids)
    if "FDR" not in sv.columns and "PValue" in sv.columns:
        _, fdr, _, _ = multipletests(sv["PValue"].values, method="fdr_bh")
        sv = sv.copy()
        sv["FDR"] = fdr

    _state["results"][name] = {"table": sv, "dtu": True}
    _state["last_result"] = name

    sig = sv[sv["FDR"] < 0.05] if "FDR" in sv.columns else sv.iloc[0:0]
    top = sv.sort_values("FDR" if "FDR" in sv.columns else "PValue").head(5)

    lines = [
        f"DTU: splice_variants completed ({len(sv):,} genes).",
        f"Genes with splice-variant signal (FDR < 0.05): {len(sig):,}",
        "",
        "Top 5 genes:",
        f"{'GeneID':<24} {'NExons':>8} {'PValue':>12} {'FDR':>12}",
        "-" * 62,
    ]
    for _, row in top.iterrows():
        gid = str(row["GeneID"])
        if len(gid) > 23:
            gid = gid[:21] + ".."
        fdr_val = float(row["FDR"]) if "FDR" in row.index else np.nan
        lines.append(
            f"{gid:<24} {int(row['NExons']):>8d} "
            f"{float(row['PValue']):>12.2e} {fdr_val:>12.2e}"
        )
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 8g: aggregate_to_genes
# ---------------------------------------------------------------------------

@mcp.tool()
def aggregate_to_genes(
    name: Optional[str] = None,
    tx2gene_path: Optional[str] = None,
    gene_column: Optional[str] = None,
    weights: str = "sleuth",
    result_name: Optional[str] = None,
) -> str:
    """Gene-level p-values from a transcript-level test, by Lancaster's method.

    Combines each gene's transcript p-values (Yi et al. 2018), so genes are
    found whether their transcripts change together or switch isoforms.
    Run after a transcript-level test (test_contrast, test_coef,
    glm_lrt_test or exact_test, e.g. on load_quant output). The gene-level
    result is stored for get_top_genes and get_result_table.

    Args:
        name: Transcript-level result to aggregate. Default: most recent.
        tx2gene_path: TSV/CSV mapping transcript IDs to gene IDs, such as kb
            ref's t2g.txt. Not needed for Gencode transcript names (their
            EnsG column is used) or when gene_column is given.
        gene_column: Gene annotation column holding gene IDs.
        weights: 'sleuth' (default; mean log normalized count of each
            transcript, as in sleuth) or 'equal' (Fisher's method).
        result_name: Name for the gene-level result. Default: '<name>_genes'.
    """
    if name is None:
        name = _state["last_result"]
    if name is None or name not in _state["results"]:
        raise ValueError("No transcript-level test result. Run test_contrast() first.")
    res = _state["results"][name]
    if "gene.table" in res or res.get("gene_sets") or res.get("dtu"):
        raise ValueError(f"Result '{name}' is not a transcript-level test.")
    if weights not in ("sleuth", "equal"):
        raise ValueError("weights must be 'sleuth' or 'equal'.")
    tx2gene = None
    if tx2gene_path is not None:
        tx2gene = pd.read_csv(tx2gene_path) if tx2gene_path.endswith(".csv") else tx2gene_path
    table = ep.aggregate_tx_pvalues(
        res, y=_state["dgelist"] if weights == "sleuth" else None,
        gene=gene_column, tx2gene=tx2gene,
        weights="equal" if weights == "equal" else None)

    out_name = result_name or f"{name}_genes"
    _state["results"][out_name] = {"gene.table": table, "table": table,
                                   "method": f"Lancaster ({weights} weights)"}
    _state["last_result"] = out_name

    sig = table["FDR"] < 0.05
    top = table.sort_values("PValue", kind="stable").head(5)
    lines = [
        f"Aggregated {int(table['NTx'].sum()):,} transcripts to {len(table):,} genes "
        f"(Lancaster, {weights} weights) as result '{out_name}'.",
        f"Genes with FDR < 0.05: {int(sig.sum()):,}",
        "",
        "Top 5 genes:",
        f"{'Gene':<20} {'NTx':>4} {'PValue':>10} {'FDR':>10}  TopTx",
        "-" * 64,
    ]
    for _, row in top.iterrows():
        name_ = row.get("GeneName")
        gene = str(name_) if isinstance(name_, str) and name_ else str(row["GeneID"])
        lines.append(f"{gene[:19]:<20} {int(row['NTx']):>4d} {row['PValue']:>10.2e} "
                     f"{row['FDR']:>10.2e}  {row['TopTx']}")
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 9: get_top_genes
# ---------------------------------------------------------------------------

def _result_table(res: dict, n: Optional[int] = None) -> pd.DataFrame:
    """Result table ranked by p-value: top_tags for DE and methylation
    tests, the gene-level table for DTU tests."""
    if "gene.table" in res:
        table = res["gene.table"]
    elif res.get("dtu"):
        table = res["table"]
    else:
        return ep.top_tags(res, n=res["table"].shape[0] if n is None else n)["table"]
    table = table.sort_values("PValue", kind="stable")
    return table if n is None else table.head(n)


@mcp.tool()
def get_top_genes(
    name: Optional[str] = None,
    n: int = 20,
    fdr_threshold: float = 0.05,
) -> str:
    """Get the top differentially expressed genes from a test result.

    Args:
        name: Which result set to query. Default: most recent.
        n: Number of top genes to return. Default: 20.
        fdr_threshold: FDR cutoff for significance. Default: 0.05.
    """
    if name is None:
        name = _state["last_result"]
    if name is None or name not in _state["results"]:
        available = list(_state["results"].keys())
        if not available:
            return "No test results available. Run test_contrast() first."
        return f"Result '{name}' not found. Available: {', '.join(available)}"

    res = _state["results"][name]
    if res.get("gene_sets"):
        return f"Result: {name}\n\n" + res["table"].head(n).to_string()
    table = _result_table(res, n)

    # Filter by FDR
    sig = table[table["FDR"] < fdr_threshold]

    lines = [
        f"Result: {name} (showing top {min(n, len(table))} genes, "
        f"{len(sig)} with FDR < {fdr_threshold})",
        "",
    ]

    # Columns present depend on the test (logFC or logOR; F, LR or none)
    value_cols = [c for c in ("logFC", "logOR", "logCPM", "F", "LR", "NExons", "NTx",
                              "TopTx.logFC") if c in table.columns]
    header = f"{'Gene':<20}" + "".join(f" {c:>8}" for c in value_cols)
    header += f" {'PValue':>10} {'FDR':>10} {'Sig':>3}"
    lines.append(header)
    lines.append("-" * len(header))

    for _, row in table.iterrows():
        gene = _gene_name(row)
        if len(gene) > 19:
            gene = gene[:17] + ".."
        sig_mark = "***" if row["FDR"] < 0.001 else "**" if row["FDR"] < 0.01 else "*" if row["FDR"] < fdr_threshold else ""
        line = f"{gene:<20}"
        for c in value_cols:
            line += f" {int(row[c]):>8d}" if c in ("NExons", "NTx") else f" {row[c]:>8.3f}"
        line += f" {row['PValue']:>10.2e} {row['FDR']:>10.2e} {sig_mark:>3}"
        lines.append(line)

    _append_gene_level_warning(lines)
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 9b: get_result_table
# ---------------------------------------------------------------------------

@mcp.tool()
def get_result_table(
    name: Optional[str] = None,
    format: str = "tsv",
    max_rows: Optional[int] = None,
) -> str:
    """Return a full result table as TSV/CSV/JSON text.

    Args:
        name: Which result set to query. Default: most recent.
        format: 'tsv', 'csv', or 'json'.
        max_rows: Optional row cap.
    """
    if name is None:
        name = _state["last_result"]
    if name is None or name not in _state["results"]:
        available = list(_state["results"].keys())
        if not available:
            return "No test results available."
        return f"Result '{name}' not found. Available: {', '.join(available)}"

    res = _state["results"][name]
    if res.get("gene_sets"):
        table = res["table"].rename_axis("Set").reset_index()
    else:
        table = _result_table(res)
    if max_rows is not None:
        table = table.iloc[:max_rows]

    if format == "json":
        return table.to_json(orient="records")
    if format == "csv":
        return table.to_csv(index=False)
    if format == "tsv":
        return table.to_csv(sep="\t", index=False)
    return "Invalid format. Use 'tsv', 'csv', or 'json'."


# ---------------------------------------------------------------------------
# Tool 9c: save_results
# ---------------------------------------------------------------------------

@mcp.tool()
def save_results(
    output_path: str,
    name: Optional[str] = None,
    format: str = "tsv",
) -> str:
    """Save a full result table to disk.

    Args:
        output_path: Destination file path.
        name: Which result set to save. Default: most recent.
        format: 'tsv', 'csv', or 'json'.
    """
    content = get_result_table(name=name, format=format)
    if content.startswith("No test results") or content.startswith("Result '"):
        return content

    with open(output_path, "w", encoding="utf-8") as f:
        f.write(content)
    return f"Saved results to: {output_path}"


# ---------------------------------------------------------------------------
# Tool 10: generate_plot
# ---------------------------------------------------------------------------

@mcp.tool()
def generate_plot(
    plot_type: str,
    result_name: Optional[str] = None,
    output_path: Optional[str] = None,
) -> str:
    """Generate and save a visualization.

    Args:
        plot_type: Type of plot. One of:
            'mds' — multi-dimensional scaling of samples,
            'bcv' — biological coefficient of variation vs abundance,
            'ql_dispersion' — quasi-likelihood dispersion plot,
            'md' — mean-difference (MA) plot for a test result,
            'volcano' — volcano plot (logFC vs -log10 p-value),
            'heatmap' — heatmap of top DE genes.
        result_name: Which test result to plot (for md/volcano/heatmap).
            Default: most recent result.
        output_path: Path to save the PNG. Default: auto-generated
            in the current working directory.
    """
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    valid_types = {"mds", "bcv", "ql_dispersion", "md", "volcano", "heatmap"}
    if plot_type not in valid_types:
        return f"Invalid plot_type '{plot_type}'. Choose from: {', '.join(sorted(valid_types))}"

    if output_path is None:
        suffix = f"_{result_name}" if result_name else ""
        output_path = os.path.join(os.getcwd(), f"{plot_type}{suffix}.png")

    d = _state["dgelist"]

    if plot_type == "mds":
        _require("dgelist", "DGEList")
        fig, ax = ep.plot_mds(d)
        fig.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        return f"MDS plot saved to: {output_path}"

    elif plot_type == "bcv":
        _require("dgelist", "DGEList")
        if not _state["dispersions_estimated"]:
            return "Dispersions not estimated yet. Run estimate_dispersion() first."
        fig, ax = ep.plot_bcv(d)
        fig.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        return f"BCV plot saved to: {output_path}"

    elif plot_type == "ql_dispersion":
        _require("fit", "fitted model")
        fig, ax = ep.plot_ql_disp(_state["fit"])
        fig.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        return f"QL dispersion plot saved to: {output_path}"

    elif plot_type == "md":
        if result_name is None:
            result_name = _state["last_result"]
        if result_name is None or result_name not in _state["results"]:
            return "No test result available. Run test_contrast() first."
        res = _state["results"][result_name]
        status = ep.decide_tests(res)
        fig, ax = ep.plot_md(res, status=status)
        ax.set_title(f"MD plot: {result_name}")
        fig.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        return f"MD plot saved to: {output_path}"

    elif plot_type == "volcano":
        if result_name is None:
            result_name = _state["last_result"]
        if result_name is None or result_name not in _state["results"]:
            return "No test result available. Run test_contrast() first."
        res = _state["results"][result_name]
        tt = ep.top_tags(res, n=res["table"].shape[0])
        table = tt["table"]

        fig, ax = plt.subplots(figsize=(8, 6))
        logfc = table["logFC"].values
        pval = table["PValue"].values
        neg_log_p = -np.log10(np.clip(pval, 1e-300, 1))
        fdr = table["FDR"].values

        # Color by significance
        colors = np.where(
            (fdr < 0.05) & (logfc > 0), "red",
            np.where((fdr < 0.05) & (logfc < 0), "blue", "grey")
        )
        ax.scatter(logfc, neg_log_p, c=colors, s=4, alpha=0.5)
        ax.set_xlabel("log2 fold change")
        ax.set_ylabel("-log10(p-value)")
        ax.set_title(f"Volcano: {result_name}")
        ax.axhline(-np.log10(0.05), color="grey", linestyle="--", linewidth=0.5)
        fig.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        return f"Volcano plot saved to: {output_path}"

    elif plot_type == "heatmap":
        if result_name is None:
            result_name = _state["last_result"]
        if result_name is None or result_name not in _state["results"]:
            return "No test result available. Run test_contrast() first."
        res = _state["results"][result_name]
        tt = ep.top_tags(res, n=30)
        table = tt["table"]

        # Get log-CPM for top genes
        logcpm = ep.cpm(d, log=True)

        # The table from top_tags has gene annotation columns (e.g. GeneID)
        # plus stat columns. The original row positions in the result table
        # correspond to gene indices in the count matrix.
        # Use the original table's row ordering from the test result.
        all_tt = ep.top_tags(res, n=res["table"].shape[0])
        all_genes = [_gene_name(row) for _, row in all_tt["table"].iterrows()]
        top_genes = [_gene_name(row) for _, row in table.iterrows()]

        gene_ids = d.get("_gene_ids", [])
        idx = []
        gene_labels = []
        for gn in top_genes:
            if gn in gene_ids:
                idx.append(gene_ids.index(gn))
                gene_labels.append(gn)
        if not idx:
            # Fallback: use row positions from the sorted result
            idx = list(range(min(30, logcpm.shape[0])))
            gene_labels = [str(i) for i in idx]

        mat = logcpm[idx, :]
        # Z-score per gene
        mat_z = (mat - mat.mean(axis=1, keepdims=True)) / (mat.std(axis=1, keepdims=True) + 1e-10)

        fig, ax = plt.subplots(figsize=(max(6, len(d.get("_sample_names", [])) * 0.6), max(6, len(idx) * 0.3)))
        im = ax.imshow(mat_z, aspect="auto", cmap="RdBu_r", vmin=-3, vmax=3)
        sample_names = d.get("_sample_names", [f"S{i+1}" for i in range(mat.shape[1])])
        ax.set_xticks(range(len(sample_names)))
        ax.set_xticklabels(sample_names, rotation=45, ha="right", fontsize=8)
        ax.set_yticks(range(len(gene_labels)))
        ax.set_yticklabels(gene_labels, fontsize=7)
        ax.set_title(f"Top DE genes: {result_name}")
        fig.colorbar(im, ax=ax, label="z-score (log-CPM)", shrink=0.6)
        fig.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        return f"Heatmap saved to: {output_path}"


# ---------------------------------------------------------------------------
# Tool 11: gene_set_test
# ---------------------------------------------------------------------------

def _read_gene_sets(gene_sets: Optional[dict], gmt_path: Optional[str]) -> dict:
    sets = {}
    if gmt_path is not None:
        with open(gmt_path, encoding="utf-8") as f:
            for line in f:
                fields = line.rstrip("\n").split("\t")
                if len(fields) >= 3:
                    sets[fields[0]] = [g for g in fields[2:] if g]
    if gene_sets:
        sets.update({str(k): list(v) for k, v in gene_sets.items()})
    if not sets:
        raise ValueError("Provide gene_sets or gmt_path.")
    return sets


@mcp.tool()
def gene_set_test(
    contrast: Optional[str] = None,
    coef: Optional[int] = None,
    gene_sets: Optional[dict] = None,
    gmt_path: Optional[str] = None,
    method: str = "camera",
    n: int = 20,
) -> str:
    """Test gene sets for differential expression with camera or fry.

    Uses the voom result if voom_transform() has been run, otherwise the
    DGEList with its dispersions (run estimate_dispersion() first).

    Args:
        contrast: Contrast string, e.g. 'groupB - groupA'.
        coef: 0-based coefficient index (alternative to contrast). Default:
            the last coefficient.
        gene_sets: Dictionary mapping set names to lists of gene IDs.
        gmt_path: Path to a GMT file of gene sets (alternative to gene_sets).
        method: 'camera' (competitive, default) or 'fry' (self-contained).
        n: Number of top sets to report.
    """
    _require("dgelist", "DGEList")
    _require("design", "design matrix")
    d = _state["dgelist"]
    if method not in ("camera", "fry"):
        raise ValueError("method must be 'camera' or 'fry'.")

    if _state["voom"] is not None:
        y, source = _state["voom"]["EList"], "voom"
    elif _state["dispersions_estimated"]:
        y, source = d, "DGEList"
    else:
        raise ValueError("Run voom_transform() or estimate_dispersion() first.")

    ids = d.get("_gene_ids")
    if ids is None:
        ids = [str(i) for i in range(d["counts"].shape[0])]
    pos = {}
    for i, g in enumerate(ids):
        pos.setdefault(str(g), i)
    index = {}
    for name, genes in _read_gene_sets(gene_sets, gmt_path).items():
        idx = sorted({pos[str(g)] for g in genes if str(g) in pos})
        if idx:
            index[name] = np.asarray(idx, dtype=int)
    if not index:
        raise ValueError("No gene set has any genes in the data.")

    design = np.asarray(_state["design"], dtype=np.float64)
    if contrast is not None:
        con = _parse_contrast(contrast, _state["design_info"])
        label = contrast
    else:
        con = design.shape[1] - 1 if coef is None else int(coef)
        label = _state["design_info"][con]

    fn = ep.camera if method == "camera" else ep.fry
    table = fn(y, index, design=design, contrast=con)
    name = f"{method}:{label.replace(' ', '')}"
    _state["results"][name] = {"table": table, "gene_sets": True}
    _state["last_result"] = name

    lines = [
        f"{method} on {source}: {label}",
        f"Sets tested: {len(index)} (FDR < 0.05: {int((table['FDR'] < 0.05).sum())})",
        "",
        f"{'Set':<30} {'NGenes':>7} {'Dir':>5} {'PValue':>10} {'FDR':>10}",
        "-" * 66,
    ]
    for set_name, row in table.head(n).iterrows():
        sn = str(set_name)
        if len(sn) > 29:
            sn = sn[:27] + ".."
        lines.append(
            f"{sn:<30} {int(row['NGenes']):>7d} {str(row.get('Direction', '')):>5} "
            f"{row['PValue']:>10.2e} {row['FDR']:>10.2e}"
        )
    return "\n".join(lines)


# ===========================================================================
# Methylation (paired-count binomial) tools
# ===========================================================================

_bin_state: dict = {
    "pclist": None,     # PCList of methylated / unmethylated counts
    "design": None,
    "design_info": None,
    "fit": None,        # bin_ql_fit result
}


@mcp.tool()
def load_bismark(
    files: list,
    sample_names: Optional[list] = None,
    group: Optional[list] = None,
    min_coverage: int = 8,
) -> str:
    """Load Bismark coverage files as paired methylated/unmethylated counts.

    Loci are kept when every sample has at least min_coverage reads, as in
    the edgeR methylation workflow (Chen et al. 2018).

    Args:
        files: Bismark .cov or .cov.gz files, one per sample.
        sample_names: Optional sample names. Default: file names.
        group: Optional group for each sample.
        min_coverage: Minimum total reads per locus in every sample.
            Default 8; 0 keeps all loci.
    """
    y = ep.read_bismark2pc(files, sample_names=sample_names, verbose=False)
    n_loci = y["counts"].shape[0]
    cov = np.asarray(y["counts"]) + np.asarray(y["counts2"])
    keep = np.all(cov >= min_coverage, axis=1)
    if not keep.all():
        y = y[np.flatnonzero(keep), :]
    if group is not None:
        if len(group) != y["counts"].shape[1]:
            raise ValueError(f"group has {len(group)} values for {y['counts'].shape[1]} samples.")
        y["samples"]["group"] = pd.Categorical([str(g) for g in group])

    _bin_state.update(pclist=y, design=None, design_info=None, fit=None)
    meth = np.asarray(y["counts"], dtype=float)
    total = meth + np.asarray(y["counts2"], dtype=float)
    beta = meth.sum(axis=0) / np.maximum(total.sum(axis=0), 1)
    names = [str(x) for x in y["samples"].index]
    lines = [
        f"Loaded {n_loci:,} loci × {len(names)} samples; "
        f"{int(keep.sum()):,} kept with coverage ≥ {min_coverage} in every sample",
        "Overall methylation: " + ", ".join(f"{n} {b:.3f}" for n, b in zip(names, beta)),
    ]
    return "\n".join(lines)


@mcp.tool()
def fit_bin_model(
    formula: Optional[str] = None,
    matrix: Optional[list] = None,
    columns: Optional[list] = None,
    robust: bool = True,
) -> str:
    """Fit quasi-binomial GLMs (edgeR 5.0's binQLFit) to the loaded loci.

    Give either an R-style formula on the sample metadata (e.g. '~group')
    or an explicit design matrix with its column names.

    Args:
        formula: Design formula. Default '~group' when neither is given.
        matrix: Design matrix as a list of rows (samples × coefficients).
        columns: Column names for matrix.
        robust: Robust empirical Bayes for the quasi-dispersions. Default True.
    """
    y = _bin_state["pclist"]
    if y is None:
        raise ValueError("No methylation data loaded. Use load_bismark() first.")
    if matrix is not None:
        design = np.asarray(matrix, dtype=np.float64)
        col_names = list(columns) if columns is not None else [f"coef{i}" for i in range(design.shape[1])]
    else:
        import patsy
        formula = formula or "~group"
        dinfo = patsy.dmatrix(formula, data=y["samples"], return_type="dataframe")
        design = np.asarray(dinfo, dtype=np.float64)
        col_names = list(dinfo.columns)
    if design.shape[0] != y["counts"].shape[1]:
        raise ValueError(f"Design has {design.shape[0]} rows for {y['counts'].shape[1]} samples.")

    fit = ep.bin_ql_fit(y, design=design, robust=robust)
    _bin_state.update(design=design, design_info=col_names, fit=fit)
    pdf = np.atleast_1d(fit.get("df.prior"))
    return "\n".join([
        f"Fitted quasi-binomial GLM: {fit['coefficients'].shape[0]:,} loci, {design.shape[1]} coefficients",
        f"Prior df (median): {np.median(pdf):.2f}",
        f"Coefficients: {', '.join(col_names)}",
    ])


@mcp.tool()
def test_bin(
    contrast: Optional[str] = None,
    coef: Optional[int] = None,
    lor: Optional[float] = None,
    name: Optional[str] = None,
) -> str:
    """Test for differential methylation (edgeR 5.0's binQLFTest).

    Results are stored with the other test results, so get_top_genes,
    get_result_table and save_results work on them.

    Args:
        contrast: Contrast string on the design columns, e.g. 'groupB'.
        coef: 0-based coefficient index (alternative to contrast). Default:
            the last coefficient.
        lor: Optional log2 odds-ratio threshold (TREAT/UPSHOT test).
        name: Label for this result set.
    """
    fit = _bin_state["fit"]
    if fit is None:
        raise ValueError("No methylation model fitted. Use fit_bin_model() first.")
    if contrast is not None:
        res = ep.bin_ql_ftest(fit, contrast=_parse_contrast(contrast, _bin_state["design_info"]), lor=lor)
        label = contrast
    else:
        c = fit["coefficients"].shape[1] - 1 if coef is None else int(coef)
        res = ep.bin_ql_ftest(fit, coef=c, lor=lor)
        label = _bin_state["design_info"][c]
    if name is None:
        name = "meth:" + label.replace(" ", "")
    _state["results"][name] = res
    _state["last_result"] = name

    table = ep.top_tags(res, n=res["table"].shape[0])["table"]
    fdr = table["FDR"].to_numpy(dtype=float)
    lor_vals = table["logOR"].to_numpy(dtype=float)
    return "\n".join([
        f"Test: binomial QL F-test — {label}" + (f" (|logOR| > {lor})" if lor else ""),
        f"Differentially methylated loci (FDR < 0.05): "
        f"{int(((fdr < 0.05) & (lor_vals > 0)).sum())} up, "
        f"{int(((fdr < 0.05) & (lor_vals < 0)).sum())} down",
        f"Result name: {name} (use get_top_genes for the top loci)",
    ])


# ===========================================================================
# Single-cell tools
# ===========================================================================

_sc_state: dict = {
    "counts": None,       # genes × cells ndarray
    "design": None,       # cells × predictors DataFrame
    "sample": None,       # per-cell sample IDs (array)
    "gene_names": None,   # gene name array
    "fit": None,          # glm_sc_fit result dict
    "n_cells": 0,
    "n_genes": 0,
    "n_samples": 0,
}


# ---------------------------------------------------------------------------
# SC Tool 1: load_sc_data
# ---------------------------------------------------------------------------

def _h5_strings(node) -> np.ndarray:
    """Decode an h5ad string or categorical array: a plain dataset, or the
    group encodings of anndata >= 0.8 (categorical, nullable-string-array)."""
    import h5py
    if isinstance(node, h5py.Dataset):
        return np.array([v.decode() if isinstance(v, bytes) else v for v in node[:]], dtype=object)
    if "codes" in node and "categories" in node:
        cats = _h5_strings(node["categories"])
        codes = node["codes"][:]
        out = np.empty(len(codes), dtype=object)
        ok = codes >= 0
        out[ok] = cats[codes[ok]]
        out[~ok] = None
        return out
    if "values" in node:
        return _h5_strings(node["values"])
    raise ValueError(f"Unsupported h5ad array encoding: {dict(node.attrs)}")


def _h5_index(grp) -> np.ndarray:
    """Row names of an h5ad obs/var dataframe group."""
    key = grp.attrs.get("_index")
    if isinstance(key, bytes):
        key = key.decode()
    if key is None or key not in grp:
        key = "index" if "index" in grp else "_index"
    return _h5_strings(grp[key])


@mcp.tool()
def load_sc_data(
    h5ad_path: str,
    sample_col: str,
    cell_type_col: Optional[str] = None,
    cell_type_filter: Optional[list] = None,
    layer: str = "raw",
    log1p_transform: bool = True,
) -> str:
    """Load single-cell count data from an h5ad file.

    Reads counts, cell metadata, and gene names. Optionally subsets
    to specific cell types.

    Args:
        h5ad_path: Path to .h5ad file.
        sample_col: Column in obs for sample/subject IDs (e.g. 'orgID').
        cell_type_col: Optional column in obs for cell type annotations.
        cell_type_filter: Optional list of cell type values to keep.
            Only used if cell_type_col is provided.
        layer: Which count layer to use — 'raw' (from adata.raw.X),
            'X' (from adata.X), or a named layer. Default: 'raw'.
        log1p_transform: If True, reverse log1p transformation on
            the counts (i.e. round(expm1(x))). Default: True.
    """
    import h5py
    import scipy.sparse as sp

    h5 = h5py.File(h5ad_path, "r")

    # --- Read counts ---
    use_raw = layer == "raw" and "raw" in h5
    if use_raw:
        x_node, var_grp = h5["raw"]["X"], h5["raw"]["var"]
    elif layer in ("X", "raw"):
        x_node, var_grp = h5["X"], h5["var"]
    else:
        if "layers" not in h5 or layer not in h5["layers"]:
            raise ValueError(f"Layer '{layer}' not found in {h5ad_path}.")
        x_node, var_grp = h5["layers"][layer], h5["var"]
    gene_names = _h5_index(var_grp)

    if isinstance(x_node, h5py.Dataset):
        data = x_node[:]
        if log1p_transform:
            data = np.round(np.expm1(data))
        counts_csr = None
        counts_dense = data.astype(np.float64)
    else:
        data = x_node["data"][:]
        if log1p_transform:
            data = np.round(np.expm1(data))
        data = data.astype(np.float64)
        indices = x_node["indices"][:]
        indptr = x_node["indptr"][:]
        enc = x_node.attrs.get("encoding-type", "csr_matrix")
        if isinstance(enc, bytes):
            enc = enc.decode()
        if enc == "csc_matrix":
            n_obs = int(x_node.attrs["shape"][0]) if "shape" in x_node.attrs else len(_h5_index(h5["obs"]))
            counts_csr = sp.csc_matrix((data, indices, indptr),
                                       shape=(n_obs, len(gene_names))).tocsr()
        else:
            counts_csr = sp.csr_matrix((data, indices, indptr),
                                       shape=(len(indptr) - 1, len(gene_names)))
        counts_dense = None

    # --- Read obs metadata ---
    obs_keys = [k for k in h5["obs"].keys() if k != "__categories"]
    categories = {}
    if "__categories" in h5["obs"]:
        for cat_name in h5["obs"]["__categories"]:
            raw_cats = h5["obs"]["__categories"][cat_name][:]
            categories[cat_name] = [c.decode() if isinstance(c, bytes) else c
                                    for c in raw_cats]

    def _decode_col(name):
        node = h5["obs"][name]
        if isinstance(node, h5py.Group):
            return _h5_strings(node)
        codes = node[:]
        if name in categories:
            return np.array([categories[name][c] for c in codes])
        if codes.dtype.kind in ("S", "O"):
            return np.array([c.decode() if isinstance(c, bytes) else c
                             for c in codes])
        return codes

    sample_ids = _decode_col(sample_col)

    # --- Cell type filtering ---
    if cell_type_col is not None and cell_type_filter is not None:
        ct_vals = _decode_col(cell_type_col)
        mask = np.isin(ct_vals, cell_type_filter)
        if counts_csr is not None:
            counts_csr = counts_csr[mask, :]
        else:
            counts_dense = counts_dense[mask, :]
        sample_ids = sample_ids[mask]
        n_obs = int(mask.sum())
        ct_counts = {ct: int((ct_vals[mask] == ct).sum()) for ct in cell_type_filter}
    else:
        if counts_csr is not None:
            n_obs = counts_csr.shape[0]
        else:
            n_obs = counts_dense.shape[0]
        ct_counts = None

    h5.close()

    # --- Convert to genes × cells ---
    if counts_csr is not None:
        counts_T = counts_csr.T.toarray().astype(np.float64)
    else:
        counts_T = counts_dense.T.astype(np.float64)

    n_genes = counts_T.shape[0]
    n_samples = len(set(sample_ids))

    _sc_state["counts"] = counts_T
    _sc_state["sample"] = sample_ids
    _sc_state["gene_names"] = gene_names
    _sc_state["design"] = None
    _sc_state["fit"] = None
    _sc_state["n_cells"] = n_obs
    _sc_state["n_genes"] = n_genes
    _sc_state["n_samples"] = n_samples

    lines = [
        f"Loaded {n_genes:,} genes × {n_obs:,} cells, {n_samples} samples",
    ]
    if ct_counts is not None:
        for ct, n in ct_counts.items():
            lines.append(f"  {ct}: {n:,} cells")
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# SC Tool 1b: load_kb
# ---------------------------------------------------------------------------

@mcp.tool()
def load_kb(
    kb_dirs: list,
    sample_names: Optional[list] = None,
    mode: str = "single_cell",
    layer: Optional[str] = None,
    filtered: Optional[bool] = None,
    min_counts: int = 0,
    group: Optional[list] = None,
) -> str:
    """Load kb count (kallisto | bustools) output, one run per sample.

    mode='single_cell' loads the cells for the single-cell mixed model, each
    run's cells labelled with its sample name (then set_sc_design, with
    sample_covariates for per-sample covariates, → fit_sc_model).
    mode='pseudobulk' sums each run's cells into one column for the bulk
    workflow (filter_genes → normalize → set_design → ...).

    Args:
        kb_dirs: kb count output directories, one per sample.
        sample_names: Sample names. Default: the directory names.
        mode: 'single_cell' (default) or 'pseudobulk'.
        layer: For the nac workflow: 'mature', 'nascent' or 'ambiguous' (or
            'cell', 'nucleus', 'total' from --sum). Default: mature +
            nascent + ambiguous, as in kb's own h5ad.
        filtered: Use counts_filtered/ (True) or counts_unfiltered/ (False).
            Default: counts_filtered/ when present.
        min_counts: Keep only barcodes with at least this many counts, e.g.
            to drop empty droplets from unfiltered output. Default 0.
        group: For mode='pseudobulk': the group of each sample.
    """
    import scipy.sparse as sp

    if mode not in ("single_cell", "pseudobulk"):
        raise ValueError("mode must be 'single_cell' or 'pseudobulk'.")
    if not kb_dirs:
        raise ValueError("Provide kb_dirs.")
    if sample_names is None:
        sample_names = [os.path.basename(os.path.normpath(d)) for d in kb_dirs]
    sample_names = [str(s) for s in sample_names]
    if len(sample_names) != len(kb_dirs) or len(set(sample_names)) != len(sample_names):
        raise ValueError("sample_names must be unique, one for each of kb_dirs.")

    mats, run_lines, unfiltered, versions = [], [], [], {}
    genes = gene_names = None
    for d, name in zip(kb_dirs, sample_names):
        a = ep.read_kb(d, layer=layer, filtered=filtered, verbose=False)
        if genes is None:
            genes = list(a.var_names)
            gene_names = list(a.var["gene_name"]) if "gene_name" in a.var else None
        elif list(a.var_names) != genes:
            raise ValueError(f"{d} does not have the same genes as {kb_dirs[0]}; "
                             "were they run with the same index?")
        totals = np.asarray(a.X.sum(axis=1)).ravel()
        keep = totals >= min_counts
        mats.append(a.X[keep])
        run_info = a.uns.get("kb", {}).get("run_info", {})
        line = f"  {name}: {int(keep.sum()):,} cells"
        if min_counts > 0:
            line += f" (of {len(keep):,} barcodes)"
        line += f", {totals[keep].sum():,.0f} counts"
        if "p_pseudoaligned" in run_info:
            line += f", {run_info['p_pseudoaligned']}% pseudoaligned"
        run_lines.append(line)
        for key in ("kallisto_version", "index_version"):
            if key in run_info:
                versions.setdefault(key, {})[name] = run_info[key]
        if os.path.basename(ep.io._kb_dirs(d, filtered)[0]) != "counts_filtered":
            unfiltered.append(name)

    notes = []
    for key, by_sample in versions.items():
        if len(set(by_sample.values())) > 1:
            notes.append(f"Warning: samples were run with different {key}s: "
                         + ", ".join(f"{n}={v}" for n, v in by_sample.items()))
    if unfiltered and min_counts == 0:
        notes.append("Note: unfiltered barcodes, including empty droplets, for "
                     f"{', '.join(unfiltered)}; run kb count --filter, or set min_counts.")

    if mode == "pseudobulk":
        counts = np.column_stack([np.asarray(m.sum(axis=0)).ravel() for m in mats])
        genes_df = pd.DataFrame({"gene_id": genes})
        if gene_names is not None:
            genes_df["gene_name"] = gene_names
        if group is not None and len(group) != len(sample_names):
            raise ValueError(f"group has {len(group)} values for {len(sample_names)} samples.")
        dgelist = ep.make_dgelist(counts, genes=genes_df,
                                  group=None if group is None else [str(g) for g in group])
        dgelist["samples"].index = sample_names
        lines = _register_loaded(dgelist, source="kb_pseudobulk", data_input=kb_dirs)
        return "\n".join(lines[:1] + ["Cells per sample:"] + run_lines + lines[1:] + notes)

    X = sp.vstack(mats).tocsr()
    sample_ids = np.concatenate([np.repeat(name, m.shape[0]) for name, m in zip(sample_names, mats)])
    _sc_state["counts"] = X.T.toarray().astype(np.float64)
    _sc_state["sample"] = sample_ids
    _sc_state["gene_names"] = np.asarray(genes, dtype=object)
    _sc_state["design"] = None
    _sc_state["fit"] = None
    _sc_state["n_cells"] = X.shape[0]
    _sc_state["n_genes"] = len(genes)
    _sc_state["n_samples"] = len(sample_names)
    lines = [f"Loaded {len(genes):,} genes × {X.shape[0]:,} cells, {len(sample_names)} samples"]
    lines += run_lines + notes
    lines.append("Use set_sc_design(sample_covariates={...}) for per-sample covariates.")
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# SC Tool 2: set_sc_design
# ---------------------------------------------------------------------------

@mcp.tool()
def set_sc_design(
    covariates: Optional[dict] = None,
    sample_covariates: Optional[dict] = None,
) -> str:
    """Set the design matrix for single-cell analysis.

    Args:
        covariates: Dictionary mapping column names to per-cell values.
            Must include 'Intercept' (all ones). Values can be lists or
            arrays of length n_cells. Example:
            {"Intercept": [1,1,...], "fed": [1,0,1,...]}
        sample_covariates: Dictionary mapping column names to per-sample
            values, {sample: value}, expanded to every cell of the sample.
            Example: {"fed": {"s1": 1, "s2": 1, "s3": 0, "s4": 0}}. An
            Intercept column is added when no 'Intercept' is given.
    """
    n = _sc_state["n_cells"]
    if n == 0:
        raise ValueError("No SC data loaded. Use load_sc_data() first.")
    if covariates is None and sample_covariates is None:
        raise ValueError("Provide covariates or sample_covariates.")

    arrays = {}
    for col, values in (covariates or {}).items():
        arr = np.asarray(values, dtype=np.float64)
        if len(arr) != n:
            raise ValueError(
                f"Column '{col}' has {len(arr)} values but expected {n} cells."
            )
        arrays[col] = arr
    if sample_covariates:
        sample = np.asarray(_sc_state["sample"]).astype(str)
        for col, mapping in sample_covariates.items():
            mapping = {str(k): v for k, v in mapping.items()}
            missing = sorted(set(sample) - set(mapping))
            if missing:
                raise ValueError(f"sample_covariates['{col}'] has no value for: {', '.join(missing)}")
            arrays[col] = np.array([mapping[s] for s in sample], dtype=np.float64)
        if "Intercept" not in arrays:
            arrays = {"Intercept": np.ones(n), **arrays}
    cols = list(arrays)

    design = pd.DataFrame(arrays, columns=cols)
    _sc_state["design"] = design
    _sc_state["fit"] = None

    return (
        f"Design matrix set: {n:,} cells × {len(cols)} predictors\n"
        f"Columns: {', '.join(cols)}"
    )


# ---------------------------------------------------------------------------
# SC Tool 3: fit_sc_model
# ---------------------------------------------------------------------------

@mcp.tool()
def fit_sc_model(
    norm_method: str = "TMM",
    verbose: bool = True,
) -> str:
    """Fit NEBULA-LN single-cell mixed model.

    Fits a per-gene negative binomial gamma mixed model with a
    sample-level random intercept. Requires load_sc_data() and
    set_sc_design() to have been called first.

    Args:
        norm_method: Normalization method — 'TMM' (default) or 'none'.
        verbose: Print progress messages.
    """
    counts = _sc_state["counts"]
    if counts is None:
        raise ValueError("No SC data loaded. Use load_sc_data() first.")
    design = _sc_state["design"]
    if design is None:
        raise ValueError("No design matrix set. Use set_sc_design() first.")
    sample = _sc_state["sample"]
    gene_names = _sc_state["gene_names"]

    import time
    t0 = time.perf_counter()

    fit = ep.glm_sc_fit(
        counts,
        design=design,
        sample=sample,
        norm_method=norm_method,
        verbose=verbose,
    )
    elapsed = time.perf_counter() - t0

    # Attach gene names
    gene_mask = fit["gene_mask"]
    if fit["genes"] is None and gene_names is not None:
        fit["genes"] = pd.DataFrame({"gene": gene_names[gene_mask]})

    _sc_state["fit"] = fit

    n_tested = fit["coefficients"].shape[0]
    n_conv = int((fit["convergence"] == 1).sum())

    return (
        f"Fitted NEBULA-LN: {n_tested:,} genes in {elapsed:.1f}s\n"
        f"Converged: {n_conv:,}/{n_tested:,}\n"
        f"Predictors: {', '.join(fit['predictor_names'])}"
    )


# ---------------------------------------------------------------------------
# SC Tool 4: test_sc_coef
# ---------------------------------------------------------------------------

@mcp.tool()
def test_sc_coef(
    coef: Optional[str] = None,
    coef_index: Optional[int] = None,
    name: Optional[str] = None,
) -> str:
    """Run a Wald test on a coefficient from the single-cell model.

    Args:
        coef: Coefficient name to test (e.g. 'fed'). Default: last coefficient.
        coef_index: 0-based coefficient index. Alternative to coef.
        name: Label for this result. Default: auto-generated.
    """
    fit = _sc_state["fit"]
    if fit is None:
        raise ValueError("No SC model fitted. Use fit_sc_model() first.")

    if coef is not None:
        c = coef
    elif coef_index is not None:
        c = coef_index
    else:
        c = None  # top_tags defaults to last coefficient

    tt = ep.top_tags(fit, n=fit["coefficients"].shape[0], coef=c)
    table = tt["table"]

    fdr = table["FDR"].values
    logfc = table["logFC"].values
    valid = ~np.isnan(fdr)

    n_up = int(((fdr < 0.05) & (logfc > 0) & valid).sum())
    n_down = int(((fdr < 0.05) & (logfc < 0) & valid).sum())
    n_sig_01 = int(((fdr < 0.1) & valid).sum())

    if name is None:
        name = coef if coef is not None else f"coef{coef_index}"
    _state["results"][name] = {"table": table, "method": "nebula_ln"}
    _state["last_result"] = name

    label = coef if coef is not None else f"coefficient {coef_index}"
    lines = [
        f"Wald test: {label}",
        f"DE genes (FDR < 0.05): {n_up} up, {n_down} down ({n_up + n_down} total)",
        f"DE genes (FDR < 0.10): {n_sig_01}",
        "",
        "Top 10 genes:",
        f"{'Gene':<30} {'logFC':>8} {'SE':>8} {'PValue':>10} {'FDR':>10}",
        "-" * 70,
    ]
    for _, row in table.head(10).iterrows():
        gene = str(row.get("gene", row.name))
        if len(gene) > 29:
            gene = gene[:27] + ".."
        lines.append(
            f"{gene:<30} {row['logFC']:>8.3f} {row['SE']:>8.3f} "
            f"{row['PValue']:>10.2e} {row['FDR']:>10.2e}"
        )

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# SC Tool 5: get_sc_top_genes
# ---------------------------------------------------------------------------

@mcp.tool()
def get_sc_top_genes(
    coef: Optional[str] = None,
    coef_index: Optional[int] = None,
    n: int = 20,
    fdr_threshold: float = 0.05,
) -> str:
    """Get top DE genes from the single-cell model.

    Args:
        coef: Coefficient name to test (e.g. 'fed'). Default: last coefficient.
        coef_index: 0-based coefficient index. Alternative to coef.
        n: Number of top genes to return. Default: 20.
        fdr_threshold: FDR cutoff for marking significance. Default: 0.05.
    """
    fit = _sc_state["fit"]
    if fit is None:
        raise ValueError("No SC model fitted. Use fit_sc_model() first.")

    c = coef if coef is not None else coef_index

    tt = ep.top_tags(fit, n=n, coef=c)
    table = tt["table"]

    label = coef if coef is not None else (
        f"coefficient {coef_index}" if coef_index is not None else "last coefficient"
    )
    lines = [
        f"Top {min(n, len(table))} genes for {label}:",
        "",
        f"{'Gene':<30} {'logFC':>8} {'SE':>8} {'z':>8} {'PValue':>10} {'FDR':>10} {'Sig':>3}",
        "-" * 82,
    ]
    for _, row in table.iterrows():
        gene = str(row.get("gene", row.name))
        if len(gene) > 29:
            gene = gene[:27] + ".."
        fdr_val = row["FDR"]
        sig = "***" if fdr_val < 0.001 else "**" if fdr_val < 0.01 else \
              "*" if fdr_val < fdr_threshold else ""
        lines.append(
            f"{gene:<30} {row['logFC']:>8.3f} {row['SE']:>8.3f} "
            f"{row['z']:>8.2f} {row['PValue']:>10.2e} {fdr_val:>10.2e} {sig:>3}"
        )

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# SC Tool 6: describe_sc
# ---------------------------------------------------------------------------

@mcp.tool()
def describe_sc() -> str:
    """Report the current state of the single-cell analysis pipeline."""
    lines = []
    if _sc_state["counts"] is None:
        return "No single-cell data loaded. Use load_sc_data() first."

    lines.append(
        f"SC data: {_sc_state['n_genes']:,} genes × "
        f"{_sc_state['n_cells']:,} cells, "
        f"{_sc_state['n_samples']} samples"
    )

    if _sc_state["design"] is not None:
        cols = list(_sc_state["design"].columns)
        lines.append(f"Design: {', '.join(cols)}")
    else:
        lines.append("Design: not set")

    fit = _sc_state["fit"]
    if fit is not None:
        n_tested = fit["coefficients"].shape[0]
        n_conv = int((fit["convergence"] == 1).sum())
        lines.append(
            f"Model: NEBULA-LN, {n_tested:,} genes tested, "
            f"{n_conv:,} converged"
        )
        lines.append(f"Predictors: {', '.join(fit['predictor_names'])}")
    else:
        lines.append("Model: not fitted")

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# SC Tool 7: save_sc_results
# ---------------------------------------------------------------------------

@mcp.tool()
def save_sc_results(
    output_path: str,
    coef: Optional[str] = None,
    coef_index: Optional[int] = None,
    format: str = "csv",
) -> str:
    """Save full single-cell DE results to disk.

    Args:
        output_path: Destination file path.
        coef: Coefficient name to test and save.
        coef_index: 0-based coefficient index. Alternative to coef.
        format: 'csv' (default), 'tsv', or 'json'.
    """
    fit = _sc_state["fit"]
    if fit is None:
        raise ValueError("No SC model fitted. Use fit_sc_model() first.")

    c = coef if coef is not None else coef_index
    tt = ep.top_tags(fit, n=fit["coefficients"].shape[0], coef=c)
    table = tt["table"]

    if format == "json":
        table.to_json(output_path, orient="records")
    elif format == "tsv":
        table.to_csv(output_path, sep="\t")
    else:
        table.to_csv(output_path)

    valid_fdr = table["FDR"].dropna()
    n_sig = int((valid_fdr < 0.05).sum())

    return (
        f"Saved {len(table):,} genes to {output_path}\n"
        f"DE at FDR < 0.05: {n_sig}"
    )


# ---------------------------------------------------------------------------
# SC Tool 8: reset_sc_state
# ---------------------------------------------------------------------------

@mcp.tool()
def reset_sc_state() -> str:
    """Reset all single-cell pipeline state."""
    _sc_state["counts"] = None
    _sc_state["design"] = None
    _sc_state["sample"] = None
    _sc_state["gene_names"] = None
    _sc_state["fit"] = None
    _sc_state["n_cells"] = 0
    _sc_state["n_genes"] = 0
    _sc_state["n_samples"] = 0
    return "Single-cell state cleared."


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    mcp.run()