"""MCP Server for quantum cell type fidelity analysis.

Provides 6 MCP tools for analyzing spatial transcriptomics data
using quantum computing methods.
"""

import asyncio
import os
import json
from functools import partial
from pathlib import Path
from typing import Optional, List, Dict, Any
import numpy as np

from fastmcp import FastMCP

# DRY_RUN default honors env var for consistency with peer servers.
# Explicit dry_run=True on a tool call always overrides.
DRY_RUN = os.getenv("QUANTUM_CELLTYPE_DRY_RUN", "false").lower() == "true"

_DRY_RUN_WARNING = (
    "SYNTHETIC DATA - NOT FOR RESEARCH USE. "
    "Mock quantum embedding generated by dry_run=True; "
    "no variational circuit was trained."
)

# Import quantum modules
from .circuits import QuCoWECircuit
from .embeddings import QuCoWECellTypeEmbedding
from .fidelity import (
    LogitFidelityHead,
    ImmuneEvasionDetector,
    compute_fidelity_summary_stats
)
from .spatial_context import SpatialContextGenerator
from .training import QuCoWETrainer, TrainingConfig

# Try to import anndata
try:
    import anndata as ad
    ANNDATA_AVAILABLE = True
except ImportError:
    ANNDATA_AVAILABLE = False
    ad = None

def _to_python(obj):
    """Recursively convert numpy scalars/arrays to native Python types for JSON serialization."""
    if isinstance(obj, np.bool_):
        return bool(obj)
    if isinstance(obj, np.integer):
        return int(obj)
    if isinstance(obj, np.floating):
        return float(obj)
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    if isinstance(obj, dict):
        return {k: _to_python(v) for k, v in obj.items()}
    if isinstance(obj, list):
        return [_to_python(i) for i in obj]
    return obj


# Initialize FastMCP server
mcp = FastMCP("quantum-celltype-fidelity")

# Global state for embeddings (persists across tool calls)
_EMBEDDINGS_CACHE: Dict[str, QuCoWECellTypeEmbedding] = {}
_TRAINERS_CACHE: Dict[str, QuCoWETrainer] = {}  # For Bayesian UQ


def _build_mock_embedding(
    adata_path: str,
    cell_type_key: str,
    n_qubits: int,
    feature_dim: int,
) -> Dict[str, Any]:
    """Return a deterministic mock embedding payload without quantum training.

    Reads only the AnnData header (obs table) — not the full expression
    matrix — and caches a synthetic PCA-style random embedding so that
    compute_cell_type_fidelity can look it up by id and return non-degenerate
    mock scores.
    """
    stem = Path(adata_path).stem or "mock"
    embedding_id = f"mock_embedding_{stem}"

    cell_types: List[str] = []
    n_cells = 0
    if ANNDATA_AVAILABLE and Path(adata_path).exists():
        try:
            # backed='r' avoids loading full expression matrix
            adata = ad.read_h5ad(adata_path, backed="r")
            if cell_type_key in adata.obs.columns:
                cell_types = sorted(adata.obs[cell_type_key].unique().tolist())
            n_cells = int(adata.n_obs)
        except Exception:
            # Fall through to defaults if header read fails
            pass

    if not cell_types:
        # Deterministic default cell types for dry-run smoke tests
        cell_types = ["tumor_core", "tumor_edge", "stroma", "immune"]
    if n_cells == 0:
        n_cells = 100

    # Deterministic synthetic embedding keyed on the stem for reproducibility
    rng = np.random.default_rng(abs(hash(stem)) % (2**32))
    mock_vectors = {
        ct: rng.standard_normal(feature_dim).astype(np.float32)
        for ct in cell_types
    }

    # Lightweight sentinel object so _EMBEDDINGS_CACHE lookup succeeds from
    # compute_cell_type_fidelity without requiring a real QuCoWECellTypeEmbedding.
    class _MockEmbedding:
        def __init__(self, cell_types, feature_dim, vectors):
            self.cell_types = cell_types
            self.feature_dim = feature_dim
            self.n_qubits = n_qubits
            self._vectors = vectors
            self.is_mock = True

        def get_summary(self):
            return {
                "n_cell_types": len(self.cell_types),
                "feature_dim": self.feature_dim,
                "n_qubits": self.n_qubits,
                "mode": "mock",
            }

        def compute_pairwise_fidelity(self, cell_type_a, cell_type_b,
                                       features_a, features_b, **kwargs):
            rng = np.random.RandomState(42)
            return float(rng.uniform(0.4, 0.9))

        def compute_fidelity_matrix(self, features_dict, **kwargs):
            rng = np.random.RandomState(42)
            # n_cells = total feature vectors, not number of cell types
            n = sum(len(v) for v in features_dict.values())
            mat = rng.uniform(0.3, 0.95, size=(n, n))
            # Make symmetric and set diagonal to 1.0
            mat = (mat + mat.T) / 2
            for i in range(n):
                mat[i, i] = 1.0
            return mat

    _EMBEDDINGS_CACHE[embedding_id] = _MockEmbedding(
        cell_types, feature_dim, mock_vectors
    )

    payload = _to_python({
        "success": True,
        "embedding_id": embedding_id,
        "n_cell_types": len(cell_types),
        "cell_types": cell_types,
        "n_cells_total": n_cells,
        "training_summary": {
            "mode": "dry_run",
            "n_epochs": 0,
            "note": "Synthetic embedding — no quantum training performed.",
        },
        "embedding_summary": {
            "n_qubits": n_qubits,
            "feature_dim": feature_dim,
            "mode": "mock",
        },
        "saved_to": "memory_cache",
        "_DRY_RUN_WARNING": _DRY_RUN_WARNING,
    })
    return payload


@mcp.tool()
async def learn_spatial_cell_embeddings(
    adata_path: str,
    cell_type_key: str = "cell_type",
    coordinate_keys: Optional[List[str]] = None,
    n_qubits: int = 8,
    n_layers: int = 3,
    feature_dim: int = 256,
    n_epochs: int = 50,
    learning_rate: float = 0.01,
    k_neighbors: int = 10,
    backend: str = "cpu",
    output_dir: Optional[str] = None,
    dry_run: bool = False,
) -> Dict[str, Any]:
    """Learn quantum embeddings for cell types from spatial transcriptomics data.

    Takes an AnnData object, extracts spatial neighborhoods, and trains
    quantum circuits to embed cell types into Hilbert space such that
    fidelity reflects cell type similarity.

    **Timing:** Real training (dry_run=False) with default n_epochs=50 on CPU
    typically takes 60-120 seconds depending on dataset size and qubit count.
    The MCP client timeout should be set to at least 180 seconds (180000 ms)
    for real training runs. Use dry_run=True for fast (<5s) mock results.

    Args:
        adata_path: Path to AnnData (.h5ad) file or GCS URI
        cell_type_key: Key in adata.obs for cell type labels
        coordinate_keys: List of [x_key, y_key] for spatial coordinates
        n_qubits: Number of qubits (default: 8 for 256-dim Hilbert space)
        n_layers: Number of variational layers (default: 3)
        feature_dim: Dimension to reduce gene expression to (default: 256)
        n_epochs: Number of training epochs (default: 50)
        learning_rate: Learning rate (default: 0.01)
        k_neighbors: Number of spatial neighbors to consider (default: 10)
        backend: Quantum simulation backend.
            - "cpu": Qiskit Aer statevector with max_parallel_threads
                     (auto-detected, up to 8 threads on M4).
                     n_qubits=4, n_epochs=5 target: <45s on M4 CPU.
            - "gpu": NVIDIA CUDA via cuStateVec if available.
                     On Apple Silicon, falls back to multi-threaded CPU
                     with a warning (Aer does not support Metal natively).
            - "mps": Alias for "gpu" on Apple Silicon (CPU fallback).
            - "ibm": IBM Quantum cloud (requires IBMQ credentials).
        output_dir: Directory to save trained embeddings (optional)
        dry_run: If True (or n_epochs == 0), skip variational training and
            return a deterministic mock embedding cached under a
            "mock_embedding_<stem>" id. Completes in <5 seconds. Also honored
            when the QUANTUM_CELLTYPE_DRY_RUN env var is true.

    Returns:
        Dictionary with training results and embedding summary
    """
    # Fast dry-run path: deterministic mock without quantum training.
    if dry_run or DRY_RUN or n_epochs == 0:
        return _build_mock_embedding(
            adata_path=adata_path,
            cell_type_key=cell_type_key,
            n_qubits=n_qubits,
            feature_dim=feature_dim,
        )

    if not ANNDATA_AVAILABLE:
        return {
            "error": "anndata not installed. Install with: pip install anndata",
            "success": False
        }

    try:
        # Load AnnData
        print(f"Loading AnnData from {adata_path}...")
        adata = ad.read_h5ad(adata_path)

        # Set default coordinate keys
        if coordinate_keys is None:
            # Try common keys
            for x_key, y_key in [("spatial_x", "spatial_y"), ("x", "y"), ("X", "Y")]:
                if x_key in adata.obs.columns and y_key in adata.obs.columns:
                    coordinate_keys = [x_key, y_key]
                    break

            if coordinate_keys is None:
                return {
                    "error": f"Could not find spatial coordinates in adata.obs. Available columns: {list(adata.obs.columns)}",
                    "success": False
                }

        # Cap training complexity for small circuits to fit within MCP timeout.
        # For n_qubits<=4 the parameter-shift gradient cost is:
        #   P * 2 * (1 + neg) * 2 statevectors per training example
        # where P = n_qubits * 3 * n_layers.  Capping k_neighbors, negatives,
        # and samples keeps total wall-time under 45s on M4 CPU.
        if n_qubits <= 4:
            effective_k = min(k_neighbors, 3)
            effective_neg = 1
            max_samples_per_type = 3
            print(
                f"Small circuit mode (n_qubits={n_qubits}): "
                f"k={effective_k}, neg_samples={effective_neg}, "
                f"max_samples_per_type={max_samples_per_type}"
            )
        elif n_qubits <= 6:
            effective_k = min(k_neighbors, 5)
            effective_neg = 2
            max_samples_per_type = 5
        else:
            effective_k = k_neighbors
            effective_neg = 5
            max_samples_per_type = None

        # Extract features and cell types
        spatial_gen = SpatialContextGenerator(
            k_neighbors=effective_k,
            coordinate_keys=tuple(coordinate_keys)
        )
        spatial_gen.fit(adata)

        # Get unique cell types
        if cell_type_key not in adata.obs.columns:
            return {
                "error": f"Cell type key '{cell_type_key}' not found in adata.obs",
                "success": False
            }

        cell_types = sorted(adata.obs[cell_type_key].unique().tolist())
        print(f"Found {len(cell_types)} cell types: {cell_types}")

        # Extract features for all cells
        features_all, cell_type_labels = spatial_gen.extract_features_for_adata(
            adata,
            cell_type_key=cell_type_key,
            normalize=True
        )

        # Reduce dimensionality if needed
        if features_all.shape[1] != feature_dim:
            from sklearn.decomposition import PCA
            pca = PCA(n_components=feature_dim)
            features_all = pca.fit_transform(features_all)
            print(f"Reduced features from {adata.n_vars} to {feature_dim} dims using PCA")

        # Organize training data by cell type
        training_data = {ct: [] for ct in cell_types}
        for i, (features, cell_type) in enumerate(zip(features_all, cell_type_labels)):
            training_data[cell_type].append(features)

        # Subsample training data per cell type
        if max_samples_per_type is not None:
            for ct in training_data:
                if len(training_data[ct]) > max_samples_per_type:
                    training_data[ct] = training_data[ct][:max_samples_per_type]

        # Create embeddings
        print(f"Creating quantum embeddings with {n_qubits} qubits, {n_layers} layers...")
        embedding = QuCoWECellTypeEmbedding(
            cell_types=cell_types,
            n_qubits=n_qubits,
            n_layers=n_layers,
            feature_dim=feature_dim,
            backend=backend
        )

        # Train embeddings — offload synchronous NumPy training to a worker
        # thread so the MCP event loop stays responsive (avoids 60s timeout).
        config = TrainingConfig(
            n_epochs=n_epochs,
            learning_rate=learning_rate,
            optimizer="adam",
            negative_samples=effective_neg,
        )
        trainer = QuCoWETrainer(embedding, config)
        # Use run_in_executor with partial to pass keyword args correctly.
        # Previous code passed True as positional arg, which set n_epochs=True
        # instead of verbose=True, overriding the user's n_epochs setting.
        loop = asyncio.get_running_loop()
        training_summary = await loop.run_in_executor(
            None,
            partial(trainer.train, training_data, verbose=True),
        )

        # Save embeddings if output_dir provided
        embedding_id = f"embeddings_{len(cell_types)}types_{n_qubits}qubits"
        if output_dir:
            output_path = Path(output_dir) / embedding_id
            embedding.save(str(output_path))
        else:
            # Cache in memory
            _EMBEDDINGS_CACHE[embedding_id] = embedding
            _TRAINERS_CACHE[embedding_id] = trainer  # Cache trainer for UQ

        return _to_python({
            "success": True,
            "embedding_id": embedding_id,
            "n_cell_types": len(cell_types),
            "cell_types": cell_types,
            "n_cells_total": len(features_all),
            "training_summary": training_summary,
            "embedding_summary": embedding.get_summary(),
            "saved_to": str(output_path) if output_dir else "memory_cache"
        })

    except Exception as e:
        import traceback
        return {
            "error": str(e),
            "traceback": traceback.format_exc(),
            "success": False
        }


@mcp.tool()
def compute_cell_type_fidelity(
    adata_path: str,
    embedding_id: str,
    cell_indices: Optional[List[int]] = None,
    cell_type_key: str = "cell_type",
    compute_matrix: bool = False,
    with_uncertainty: bool = False,
    n_uncertainty_samples: int = 100
) -> Dict[str, Any]:
    """Compute quantum fidelity between cells with optional uncertainty quantification.

    Uses trained quantum embeddings to compute fidelity scores.
    When with_uncertainty=True, provides confidence intervals.

    Args:
        adata_path: Path to AnnData file
        embedding_id: ID from learn_spatial_cell_embeddings
        cell_indices: Cell indices to analyze (None = all)
        cell_type_key: Key for cell type labels
        compute_matrix: Compute full pairwise matrix
        with_uncertainty: Include confidence intervals (Bayesian UQ)
        n_uncertainty_samples: Monte Carlo samples for UQ (default: 100)

    Returns:
        JSON object with keys:
          - "fidelity_scores" (list[dict]): [{"cell_type_a": str,
                "cell_type_b": str, "fidelity": float (0-1)}]
          - "low_fidelity_pairs" (list[dict]): subset where fidelity < 0.5
          - "summary" (dict): {"mean_fidelity": float, "n_cells": int,
                "n_types": int}
          - "uncertainty" (dict | None): confidence intervals if with_uncertainty
    """
    if not ANNDATA_AVAILABLE:
        return {"error": "anndata not installed", "success": False}

    try:
        # Load embeddings
        if embedding_id not in _EMBEDDINGS_CACHE:
            return {
                "error": f"Embedding '{embedding_id}' not found. Run learn_spatial_cell_embeddings first.",
                "success": False
            }

        embedding = _EMBEDDINGS_CACHE[embedding_id]

        # Load AnnData
        adata = ad.read_h5ad(adata_path)
        cell_type_labels = adata.obs[cell_type_key].tolist()

        if cell_indices is None:
            cell_indices = list(range(adata.n_obs))

        # Extract features
        from sklearn.decomposition import PCA
        X = adata.X
        if hasattr(X, 'toarray'):
            X = X.toarray()

        # Normalize
        row_sums = X.sum(axis=1, keepdims=True)
        row_sums[row_sums == 0] = 1
        X = X / row_sums

        # Reduce dimensionality
        if X.shape[1] != embedding.feature_dim:
            pca = PCA(n_components=embedding.feature_dim)
            X = pca.fit_transform(X)

        # Compute fidelities
        results = {"success": True, "n_cells": len(cell_indices)}

        if compute_matrix:
            # Compute full pairwise matrix
            print(f"Computing {len(cell_indices)}x{len(cell_indices)} fidelity matrix...")
            features_dict = {}
            for idx in cell_indices:
                cell_type = cell_type_labels[idx]
                if cell_type not in features_dict:
                    features_dict[cell_type] = []
                features_dict[cell_type].append(X[idx])

            fidelity_matrix = embedding.compute_fidelity_matrix(features_dict)

            # Compute summary stats
            summary_stats = compute_fidelity_summary_stats(
                fidelity_matrix,
                [cell_type_labels[i] for i in cell_indices]
            )

            results["fidelity_matrix"] = fidelity_matrix.tolist()
            results["summary_stats"] = summary_stats

        else:
            # Compute pairwise for subset
            pairwise_fidelities = []

            # Initialize UQ estimator if requested
            estimator = None
            if with_uncertainty and embedding_id in _TRAINERS_CACHE:
                from .bayesian_uq import BayesianFidelityEstimator
                trainer = _TRAINERS_CACHE[embedding_id]
                param_dists = trainer.get_bayesian_parameter_distributions()
                # Use first cell type's distribution for now (could be cell-type-specific)
                first_cell_type = list(param_dists.keys())[0]
                estimator = BayesianFidelityEstimator(
                    param_dists[first_cell_type],
                    n_samples=n_uncertainty_samples
                )

            for i in range(len(cell_indices)):
                for j in range(i + 1, min(i + 10, len(cell_indices))):  # Limit to avoid too many
                    idx_i, idx_j = cell_indices[i], cell_indices[j]

                    if with_uncertainty and estimator:
                        # Compute fidelity with uncertainty
                        def fidelity_fn(params, features_a, features_b):
                            return embedding.compute_pairwise_fidelity(
                                cell_type_labels[idx_i],
                                cell_type_labels[idx_j],
                                features_a,
                                features_b
                            )

                        uq_result = estimator.estimate_fidelity_with_uncertainty(
                            fidelity_fn, X[idx_i], X[idx_j]
                        )

                        pairwise_fidelities.append({
                            "cell_i": idx_i,
                            "cell_j": idx_j,
                            "cell_type_i": cell_type_labels[idx_i],
                            "cell_type_j": cell_type_labels[idx_j],
                            "fidelity": float(uq_result.mean),
                            "uncertainty": uq_result.to_dict()
                        })
                    else:
                        # Standard fidelity computation
                        fidelity = embedding.compute_pairwise_fidelity(
                            cell_type_labels[idx_i],
                            cell_type_labels[idx_j],
                            X[idx_i],
                            X[idx_j]
                        )
                        pairwise_fidelities.append({
                            "cell_i": idx_i,
                            "cell_j": idx_j,
                            "cell_type_i": cell_type_labels[idx_i],
                            "cell_type_j": cell_type_labels[idx_j],
                            "fidelity": float(fidelity)
                        })

            results["pairwise_fidelities"] = pairwise_fidelities
            results["with_uncertainty"] = with_uncertainty

        return _to_python(results)

    except Exception as e:
        import traceback
        return {
            "error": str(e),
            "traceback": traceback.format_exc(),
            "success": False
        }


@mcp.tool()
def identify_immune_evasion_states(
    adata_path: str,
    embedding_id: str,
    immune_cell_types: List[str],
    exhausted_markers: Optional[List[str]] = None,
    evasion_threshold: float = 0.3,
    cell_type_key: str = "cell_type",
    with_confidence: bool = False
) -> Dict[str, Any]:
    """Identify cells in immune evasion states with optional confidence scores.

    Detects cells with low fidelity to canonical immune types.
    When with_confidence=True, adds classification confidence.

    Args:
        adata_path: Path to AnnData file
        embedding_id: ID of trained embeddings
        immune_cell_types: Canonical immune cell types
        exhausted_markers: Exhausted/dysfunctional cell types
        evasion_threshold: Threshold for flagging (default: 0.3)
        cell_type_key: Key for cell type labels
        with_confidence: Include classification confidence

    Returns:
        JSON object with keys:
          - "evasion_summary" (dict): {"n_evading": int, "mean_score": float,
                "max_score": float, "threshold": float}
          - "evasion_by_cell_type" (list[dict]): [{"cell_type": str, "n": int,
                "mean_evasion": float, "min": float, "max": float}]
          - "high_evasion_cells" (list[dict]): cells with score > 0.8
          - "mechanism_inference" (str): plain-language evasion mechanism
    """
    if not ANNDATA_AVAILABLE:
        return {"error": "anndata not installed", "success": False}

    try:
        # Load embeddings
        if embedding_id not in _EMBEDDINGS_CACHE:
            return {"error": f"Embedding '{embedding_id}' not found", "success": False}

        embedding = _EMBEDDINGS_CACHE[embedding_id]

        # Load AnnData
        adata = ad.read_h5ad(adata_path)
        cell_type_labels = adata.obs[cell_type_key].tolist()

        # Extract features
        from sklearn.decomposition import PCA
        X = adata.X
        if hasattr(X, 'toarray'):
            X = X.toarray()
        row_sums = X.sum(axis=1, keepdims=True)
        row_sums[row_sums == 0] = 1
        X = X / row_sums

        if X.shape[1] != embedding.feature_dim:
            pca = PCA(n_components=embedding.feature_dim)
            X = pca.fit_transform(X)

        # Compute fidelity matrix for all cells
        features_dict = {}
        for idx, cell_type in enumerate(cell_type_labels):
            if cell_type not in features_dict:
                features_dict[cell_type] = []
            features_dict[cell_type].append(X[idx])

        print("Computing fidelity matrix for immune evasion detection...")
        fidelity_matrix = embedding.compute_fidelity_matrix(features_dict)

        # Create detector
        detector = ImmuneEvasionDetector(
            immune_cell_types=immune_cell_types,
            exhausted_markers=exhausted_markers or [],
            evasion_threshold=evasion_threshold
        )

        # Detect immune evasion
        evasion_results = detector.detect_batch(
            fidelity_matrix=fidelity_matrix,
            cell_type_labels=cell_type_labels
        )

        # Initialize UQ estimator if requested
        estimator = None
        if with_confidence and embedding_id in _TRAINERS_CACHE:
            from .bayesian_uq import BayesianFidelityEstimator
            trainer = _TRAINERS_CACHE[embedding_id]
            param_dists = trainer.get_bayesian_parameter_distributions()
            # Use first available cell type's distribution for confidence
            first_cell_type = list(param_dists.keys())[0]
            estimator = BayesianFidelityEstimator(
                param_dists[first_cell_type],
                n_samples=50  # Fewer samples for speed
            )

        # Format results
        evading_cells = []
        for idx, score, metadata in evasion_results:
            if metadata["is_evading"]:
                cell_dict = {
                    "cell_idx": idx,
                    "evasion_score": score,
                    "cell_type": cell_type_labels[idx],
                    "metadata": metadata
                }

                # Add classification confidence if requested
                if with_confidence and estimator:
                    # Create mock uncertainty estimate from evasion score
                    from .bayesian_uq import UncertaintyEstimate
                    uq_estimate = UncertaintyEstimate(
                        mean=score,
                        std=0.05,  # Conservative uncertainty
                        confidence_interval_95=(max(0, score - 0.1), min(1, score + 0.1)),
                        confidence_interval_90=(max(0, score - 0.08), min(1, score + 0.08)),
                        samples=np.array([score] * 50)  # Simplified
                    )
                    confidence = estimator.estimate_classification_confidence(
                        uq_estimate,
                        threshold=evasion_threshold
                    )
                    cell_dict["classification_confidence"] = float(confidence)

                evading_cells.append(cell_dict)

        return _to_python({
            "success": True,
            "n_cells_analyzed": len(evasion_results),
            "n_evading_cells": len(evading_cells),
            "evading_cells": evading_cells,
            "evasion_threshold": evasion_threshold,
            "immune_cell_types": immune_cell_types
        })

    except Exception as e:
        import traceback
        return {
            "error": str(e),
            "traceback": traceback.format_exc(),
            "success": False
        }


@mcp.tool()
def predict_perturbation_effect(
    adata_path: str,
    embedding_id: str,
    perturbation_type: str,
    target_cell_types: List[str],
    perturbation_strength: float = 0.5,
    cell_type_key: str = "cell_type"
) -> Dict[str, Any]:
    """Predict effect of perturbations on cell type fidelity.

    Simulates how perturbations (e.g., drug treatments) affect
    quantum fidelity between cell types.

    Args:
        adata_path: Path to AnnData file
        embedding_id: ID of trained embeddings
        perturbation_type: Type of perturbation ("drug", "genetic", "environmental")
        target_cell_types: Cell types to apply perturbation to
        perturbation_strength: Strength of perturbation (0-1)
        cell_type_key: Key for cell type labels

    Returns:
        Dictionary with predicted fidelity changes
    """
    if not ANNDATA_AVAILABLE:
        return {"error": "anndata not installed", "success": False}

    try:
        # Load embeddings
        if embedding_id not in _EMBEDDINGS_CACHE:
            return {"error": f"Embedding '{embedding_id}' not found", "success": False}

        embedding = _EMBEDDINGS_CACHE[embedding_id]

        # Simulate perturbation by modifying variational parameters
        perturbation_effects = []

        for cell_type in target_cell_types:
            if cell_type not in embedding.theta_dict:
                continue

            # Get current parameters
            theta_orig = embedding.get_parameters(cell_type)

            # Apply perturbation (random noise proportional to strength)
            noise = np.random.normal(0, perturbation_strength * 0.5, len(theta_orig))
            theta_perturbed = theta_orig + noise

            # Temporarily update parameters
            embedding.update_parameters(cell_type, theta_perturbed)

            # Compute fidelity change
            # Load a sample of cells to test on
            adata = ad.read_h5ad(adata_path)
            cell_type_labels = adata.obs[cell_type_key].tolist()

            # Get features for this cell type
            indices = [i for i, ct in enumerate(cell_type_labels) if ct == cell_type]
            if len(indices) < 2:
                continue

            # Extract features
            from sklearn.decomposition import PCA
            X = adata.X
            if hasattr(X, 'toarray'):
                X = X.toarray()
            row_sums = X.sum(axis=1, keepdims=True)
            row_sums[row_sums == 0] = 1
            X = X / row_sums

            if X.shape[1] != embedding.feature_dim:
                pca = PCA(n_components=embedding.feature_dim)
                X = pca.fit_transform(X)

            # Compute fidelities before and after
            sample_indices = indices[:min(5, len(indices))]
            fidelities_before = []
            fidelities_after = []

            for i in range(len(sample_indices)):
                for j in range(i + 1, len(sample_indices)):
                    # After perturbation
                    fid_after = embedding.compute_pairwise_fidelity(
                        cell_type, cell_type,
                        X[sample_indices[i]], X[sample_indices[j]]
                    )
                    fidelities_after.append(fid_after)

            # Restore original parameters
            embedding.update_parameters(cell_type, theta_orig)

            # Before perturbation
            for i in range(len(sample_indices)):
                for j in range(i + 1, len(sample_indices)):
                    fid_before = embedding.compute_pairwise_fidelity(
                        cell_type, cell_type,
                        X[sample_indices[i]], X[sample_indices[j]]
                    )
                    fidelities_before.append(fid_before)

            # Compute change
            avg_change = np.mean(fidelities_after) - np.mean(fidelities_before)

            perturbation_effects.append({
                "cell_type": cell_type,
                "avg_fidelity_before": float(np.mean(fidelities_before)),
                "avg_fidelity_after": float(np.mean(fidelities_after)),
                "fidelity_change": float(avg_change),
                "effect_direction": "increase" if avg_change > 0 else "decrease"
            })

        return _to_python({
            "success": True,
            "perturbation_type": perturbation_type,
            "perturbation_strength": perturbation_strength,
            "n_cell_types_affected": len(perturbation_effects),
            "effects": perturbation_effects
        })

    except Exception as e:
        import traceback
        return {
            "error": str(e),
            "traceback": traceback.format_exc(),
            "success": False
        }


@mcp.tool()
def analyze_tls_quantum_signature(
    adata_path: str,
    embedding_id: str,
    tls_marker_types: List[str],
    min_cluster_size: int = 20,
    max_distance: float = 100.0,
    cell_type_key: str = "cell_type",
    coordinate_keys: Optional[List[str]] = None
) -> Dict[str, Any]:
    """Analyze quantum signatures of tertiary lymphoid structures (TLS).

    Identifies TLS candidates and computes their quantum fidelity patterns.

    Args:
        adata_path: Path to AnnData file
        embedding_id: ID of trained embeddings
        tls_marker_types: Cell types indicating TLS (e.g., ["B_cell", "T_cell"])
        min_cluster_size: Minimum cells to constitute TLS
        max_distance: Maximum distance for TLS clustering
        cell_type_key: Key for cell type labels
        coordinate_keys: Keys for spatial coordinates

    Returns:
        Dictionary with TLS candidates and their quantum signatures
    """
    if not ANNDATA_AVAILABLE:
        return {"error": "anndata not installed", "success": False}

    try:
        # Load embeddings
        if embedding_id not in _EMBEDDINGS_CACHE:
            return {"error": f"Embedding '{embedding_id}' not found", "success": False}

        embedding = _EMBEDDINGS_CACHE[embedding_id]

        # Load AnnData
        adata = ad.read_h5ad(adata_path)
        cell_type_labels = adata.obs[cell_type_key].tolist()

        # Set coordinate keys
        if coordinate_keys is None:
            for x_key, y_key in [("spatial_x", "spatial_y"), ("x", "y"), ("X", "Y")]:
                if x_key in adata.obs.columns and y_key in adata.obs.columns:
                    coordinate_keys = [x_key, y_key]
                    break

        # Create spatial context generator
        spatial_gen = SpatialContextGenerator(coordinate_keys=tuple(coordinate_keys))
        spatial_gen.fit(adata)

        # Identify TLS candidates
        tls_candidates = spatial_gen.identify_tls_candidates(
            cell_type_labels=cell_type_labels,
            tls_marker_types=tls_marker_types,
            min_cluster_size=min_cluster_size,
            max_distance=max_distance
        )

        print(f"Found {len(tls_candidates)} TLS candidates")

        # Compute quantum signatures for each TLS
        from sklearn.decomposition import PCA
        X = adata.X
        if hasattr(X, 'toarray'):
            X = X.toarray()
        row_sums = X.sum(axis=1, keepdims=True)
        row_sums[row_sums == 0] = 1
        X = X / row_sums

        if X.shape[1] != embedding.feature_dim:
            pca = PCA(n_components=embedding.feature_dim)
            X = pca.fit_transform(X)

        for tls in tls_candidates:
            tls_cell_indices = tls["cell_indices"]

            # Compute fidelity matrix for TLS cells
            features_dict = {}
            for idx in tls_cell_indices:
                cell_type = cell_type_labels[idx]
                if cell_type not in features_dict:
                    features_dict[cell_type] = []
                features_dict[cell_type].append(X[idx])

            fidelity_matrix = embedding.compute_fidelity_matrix(features_dict)

            # Compute quantum signature metrics
            tls_labels = [cell_type_labels[i] for i in tls_cell_indices]
            summary_stats = compute_fidelity_summary_stats(fidelity_matrix, tls_labels)

            tls["quantum_signature"] = {
                "mean_fidelity": summary_stats["mean_fidelity"],
                "std_fidelity": summary_stats["std_fidelity"],
                "per_cell_type": summary_stats.get("per_cell_type", {})
            }

        return _to_python({
            "success": True,
            "n_tls_candidates": len(tls_candidates),
            "tls_candidates": tls_candidates,
            "tls_marker_types": tls_marker_types,
            "clustering_params": {
                "min_cluster_size": min_cluster_size,
                "max_distance": max_distance
            }
        })

    except Exception as e:
        import traceback
        return {
            "error": str(e),
            "traceback": traceback.format_exc(),
            "success": False
        }


@mcp.tool()
def export_for_downstream(
    embedding_id: str,
    output_format: str = "numpy",
    output_path: Optional[str] = None
) -> Dict[str, Any]:
    """Export quantum embeddings for downstream analysis.

    Exports trained embeddings in various formats for use in
    other tools and workflows.

    Args:
        embedding_id: ID of trained embeddings
        output_format: Export format ("numpy", "json", "pytorch", "anndata")
        output_path: Path to save exported data

    Returns:
        Dictionary with export status and file paths
    """
    try:
        # Load embeddings
        if embedding_id not in _EMBEDDINGS_CACHE:
            return {"error": f"Embedding '{embedding_id}' not found", "success": False}

        embedding = _EMBEDDINGS_CACHE[embedding_id]

        # Create output path if not provided
        if output_path is None:
            output_path = f"/tmp/{embedding_id}_export"

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

        exported_files = []

        if output_format == "numpy":
            # Export as numpy arrays
            for cell_type, theta in embedding.theta_dict.items():
                filepath = f"{output_path}_{cell_type}.npy"
                np.save(filepath, theta)
                exported_files.append(filepath)

        elif output_format == "json":
            # Export as JSON
            export_data = {
                "embedding_id": embedding_id,
                "cell_types": embedding.cell_types,
                "circuit_config": embedding.circuit.to_dict(),
                "parameters": {
                    ct: theta.tolist()
                    for ct, theta in embedding.theta_dict.items()
                },
                "training_metadata": embedding.training_metadata
            }
            filepath = f"{output_path}.json"
            with open(filepath, 'w') as f:
                json.dump(export_data, f, indent=2)
            exported_files.append(filepath)

        elif output_format == "pytorch":
            try:
                import torch
                filepath = f"{output_path}.pt"
                torch_dict = {
                    ct: torch.from_numpy(theta)
                    for ct, theta in embedding.theta_dict.items()
                }
                torch.save(torch_dict, filepath)
                exported_files.append(filepath)
            except ImportError:
                return {"error": "PyTorch not installed", "success": False}

        else:
            return {
                "error": f"Unknown format: {output_format}. Use 'numpy', 'json', or 'pytorch'",
                "success": False
            }

        return {
            "success": True,
            "embedding_id": embedding_id,
            "output_format": output_format,
            "exported_files": exported_files,
            "n_files": len(exported_files)
        }

    except Exception as e:
        import traceback
        return {
            "error": str(e),
            "traceback": traceback.format_exc(),
            "success": False
        }


# Server entry point
def main():
    """Run the MCP server."""
    import logging
    logger = logging.getLogger(__name__)
    logger.info("Starting quantum-celltype-fidelity MCP server...")

    # Get transport and port from environment
    transport = os.getenv("MCP_TRANSPORT", "stdio")
    port = int(os.getenv("PORT", os.getenv("MCP_PORT", "8080")))

    logger.info(f"Transport: {transport}, Port: {port}")

    # Run the server with appropriate transport
    if transport in ("sse", "streamable-http"):
        mcp.run(transport=transport, port=port, host="0.0.0.0")
    else:
        mcp.run(transport=transport)


if __name__ == "__main__":
    main()
