"""
Decoupler enrichment analysis tutorial demonstrating observation-level enrichment scoring.

This MCP Server provides 4 tools:
1. decoupler_load_toy_data: Load toy dataset and visualize gene expression patterns
2. decoupler_run_individual_methods: Run individual enrichment methods (GSEA and ULM)
3. decoupler_run_multiple_methods: Run all enrichment methods and compute consensus scores
4. decoupler_access_prior_knowledge: Query OmniPath resources for prior knowledge networks

All tools extracted from `scverse/decoupler-tutorials/blob/main/example.ipynb`.
"""

import os
from datetime import datetime
from pathlib import Path

# Standard imports
from typing import Annotated

import anndata as ad
import decoupler as dc

# Domain-specific imports
import matplotlib.pyplot as plt
import pandas as pd
import scanpy as sc
from fastmcp import FastMCP

# Project structure
PROJECT_ROOT = Path(__file__).parent.parent.parent.resolve()
DEFAULT_INPUT_DIR = PROJECT_ROOT / "tmp" / "inputs"
DEFAULT_OUTPUT_DIR = PROJECT_ROOT / "tmp" / "outputs"

INPUT_DIR = Path(os.environ.get("EXAMPLE_INPUT_DIR", DEFAULT_INPUT_DIR))
OUTPUT_DIR = Path(os.environ.get("EXAMPLE_OUTPUT_DIR", DEFAULT_OUTPUT_DIR))

# Ensure directories exist
INPUT_DIR.mkdir(parents=True, exist_ok=True)
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

# Timestamp for unique outputs
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

# Configure plotting parameters
plt.rcParams["figure.dpi"] = 300
plt.rcParams["savefig.dpi"] = 300
sc.set_figure_params(figsize=(3.5, 3.5), frameon=False, dpi=300)

# MCP server instance
example_mcp = FastMCP(name="example")


@example_mcp.tool
def decoupler_load_toy_data(
    out_prefix: Annotated[str | None, "Output file prefix"] = None,
) -> dict:
    """
    Load decoupler toy dataset and visualize gene expression patterns across cell groups.
    Input is no primary data (uses built-in toy dataset) and output is AnnData object with heatmap and violin plot visualizations.
    """

    # Set matplotlib backend for non-interactive plotting
    plt.switch_backend("Agg")

    # Set output prefix
    if out_prefix is None:
        out_prefix = f"toy_data_{timestamp}"

    # Load toy dataset
    adata, net = dc.ds.toy()

    # Save the AnnData object and network
    adata_path = OUTPUT_DIR / f"{out_prefix}_adata.h5ad"
    net_path = OUTPUT_DIR / f"{out_prefix}_network.csv"

    adata.write(adata_path)
    net.to_csv(net_path, index=False)

    # Generate heatmap visualization
    heatmap_path = OUTPUT_DIR / f"{out_prefix}_gene_expression_heatmap.png"
    plt.figure(figsize=(8, 6))
    sc.pl.heatmap(adata=adata, groupby="group", var_names=adata.var_names, show=False)
    plt.savefig(heatmap_path, dpi=300, bbox_inches="tight")
    plt.close("all")

    # Generate violin plot
    violin_path = OUTPUT_DIR / f"{out_prefix}_gene_expression_violin.png"
    plt.figure(figsize=(8, 6))
    sc.pl.violin(adata, groupby="group", keys=["G01", "G05", "G13"], show=False)
    plt.savefig(violin_path, dpi=300, bbox_inches="tight")
    plt.close("all")

    # Generate network visualization
    network_path = OUTPUT_DIR / f"{out_prefix}_network_plot.png"
    plt.figure(figsize=(5, 5))
    dc.pl.network(
        net,
        size_node=15,
        s_cmap="white",
        t_cmap="white",
        c_pos_w="darkgreen",
        c_neg_w="darkred",
        return_fig=True,
    )
    plt.savefig(network_path, dpi=300, bbox_inches="tight")
    plt.close("all")

    plt.close("all")  # Close all figures to free memory

    return {
        "message": f"Toy dataset loaded and visualized with {adata.n_obs} cells and {adata.n_vars} genes",
        "reference": "https://github.com/scverse/decoupler-tutorials/blob/main/example.ipynb",
        "artifacts": [
            {
                "description": "AnnData object with expression data",
                "path": str(adata_path.resolve()),
            },
            {"description": "Prior knowledge network", "path": str(net_path.resolve())},
            {
                "description": "Gene expression heatmap",
                "path": str(heatmap_path.resolve()),
            },
            {
                "description": "Gene expression violin plot",
                "path": str(violin_path.resolve()),
            },
            {
                "description": "Network visualization",
                "path": str(network_path.resolve()),
            },
        ],
    }


@example_mcp.tool
def decoupler_run_individual_methods(
    adata_path: Annotated[str | None, "Path to AnnData file with expression data"] = None,
    network_path: Annotated[
        str | None,
        "Path to network file in CSV format with columns: source, target, weight",
    ] = None,
    tmin: Annotated[int, "Minimum number of targets per source"] = 0,
    out_prefix: Annotated[str | None, "Output file prefix"] = None,
) -> dict:
    """
    Run individual enrichment methods (GSEA and ULM) on gene expression data using prior knowledge networks.
    Input is AnnData expression file and network file and output is enrichment scores with heatmap visualizations.
    """

    # Set matplotlib backend for non-interactive plotting
    plt.switch_backend("Agg")

    # Input validation
    if adata_path is None:
        raise ValueError("Path to AnnData file must be provided")
    if network_path is None:
        raise ValueError("Path to network file must be provided")

    # File existence validation
    adata_file = Path(adata_path)
    if not adata_file.exists():
        raise FileNotFoundError(f"AnnData file not found: {adata_path}")

    network_file = Path(network_path)
    if not network_file.exists():
        raise FileNotFoundError(f"Network file not found: {network_path}")

    # Set output prefix
    if out_prefix is None:
        out_prefix = f"individual_methods_{timestamp}"

    # Load data
    adata = ad.read_h5ad(adata_path)
    net = pd.read_csv(network_path)

    # Run GSEA method
    dc.mt.gsea(data=adata, net=net, tmin=tmin)

    # Visualize GSEA scores
    scores_gsea = dc.pp.get_obsm(adata, key="score_gsea")
    gsea_heatmap_path = OUTPUT_DIR / f"{out_prefix}_gsea_scores_heatmap.png"
    plt.figure(figsize=(8, 6))
    sc.pl.heatmap(
        adata=scores_gsea,
        groupby="group",
        var_names=scores_gsea.var_names,
        cmap="RdBu_r",
        vcenter=0,
        show=False,
    )
    plt.savefig(gsea_heatmap_path, dpi=300, bbox_inches="tight")
    plt.close("all")

    # Run ULM method
    dc.mt.ulm(data=adata, net=net, tmin=tmin)

    # Visualize ULM scores
    scores_ulm = dc.pp.get_obsm(adata, key="score_ulm")
    ulm_heatmap_path = OUTPUT_DIR / f"{out_prefix}_ulm_scores_heatmap.png"
    plt.figure(figsize=(8, 6))
    sc.pl.heatmap(
        adata=scores_ulm,
        groupby="group",
        var_names=scores_ulm.var_names,
        cmap="RdBu_r",
        vcenter=0,
        show=False,
    )
    plt.savefig(ulm_heatmap_path, dpi=300, bbox_inches="tight")
    plt.close("all")

    plt.close("all")  # Close all figures to free memory

    # Save results
    adata_results_path = OUTPUT_DIR / f"{out_prefix}_adata_with_scores.h5ad"
    adata.write(adata_results_path)

    # Save score matrices as CSV
    gsea_scores_path = OUTPUT_DIR / f"{out_prefix}_gsea_scores.csv"
    scores_gsea.to_df().to_csv(gsea_scores_path)

    ulm_scores_path = OUTPUT_DIR / f"{out_prefix}_ulm_scores.csv"
    scores_ulm.to_df().to_csv(ulm_scores_path)

    return {
        "message": f"Individual enrichment methods completed: GSEA and ULM on {adata.n_obs} cells",
        "reference": "https://github.com/scverse/decoupler-tutorials/blob/main/example.ipynb",
        "artifacts": [
            {
                "description": "AnnData with enrichment scores",
                "path": str(adata_results_path.resolve()),
            },
            {
                "description": "GSEA enrichment scores",
                "path": str(gsea_scores_path.resolve()),
            },
            {
                "description": "ULM enrichment scores",
                "path": str(ulm_scores_path.resolve()),
            },
            {
                "description": "GSEA scores heatmap",
                "path": str(gsea_heatmap_path.resolve()),
            },
            {
                "description": "ULM scores heatmap",
                "path": str(ulm_heatmap_path.resolve()),
            },
        ],
    }


@example_mcp.tool
def decoupler_run_multiple_methods(
    adata_path: Annotated[str | None, "Path to AnnData file with expression data"] = None,
    network_path: Annotated[
        str | None,
        "Path to network file in CSV format with columns: source, target, weight",
    ] = None,
    methods: Annotated[str, "Methods to run - use 'all' for all available methods"] = "all",
    tmin: Annotated[int, "Minimum number of targets per source"] = 0,
    out_prefix: Annotated[str | None, "Output file prefix"] = None,
) -> dict:
    """
    Run multiple enrichment methods and compute consensus scores across methods for gene expression analysis.
    Input is AnnData expression file and network file and output is consensus enrichment scores with heatmap visualization.
    """

    # Set matplotlib backend for non-interactive plotting
    plt.switch_backend("Agg")

    # Input validation
    if adata_path is None:
        raise ValueError("Path to AnnData file must be provided")
    if network_path is None:
        raise ValueError("Path to network file must be provided")

    # File existence validation
    adata_file = Path(adata_path)
    if not adata_file.exists():
        raise FileNotFoundError(f"AnnData file not found: {adata_path}")

    network_file = Path(network_path)
    if not network_file.exists():
        raise FileNotFoundError(f"Network file not found: {network_path}")

    # Set output prefix
    if out_prefix is None:
        out_prefix = f"multiple_methods_{timestamp}"

    # Load data
    adata = ad.read_h5ad(adata_path)
    net = pd.read_csv(network_path)

    # Run multiple methods
    # Handle methods parameter - convert to list if needed to avoid system crashes
    if methods == "all":
        # Use a stable subset to prevent system crashes with all methods
        methods_to_use = ["gsea", "ulm", "mlm", "viper"]
    else:
        methods_to_use = methods

    dc.mt.decouple(
        data=adata,
        net=net,
        methods=methods_to_use,
        tmin=tmin,
    )

    # Compute consensus scores
    dc.mt.consensus(result=adata)

    # Visualize consensus scores
    scores_consensus = dc.pp.get_obsm(adata, key="score_consensus")
    consensus_heatmap_path = OUTPUT_DIR / f"{out_prefix}_consensus_scores_heatmap.png"
    plt.figure(figsize=(8, 6))
    sc.pl.heatmap(
        adata=scores_consensus,
        groupby="group",
        var_names=scores_consensus.var_names,
        cmap="RdBu_r",
        vcenter=0,
        show=False,
    )
    plt.savefig(consensus_heatmap_path, dpi=300, bbox_inches="tight")
    plt.close("all")

    plt.close("all")  # Close all figures to free memory

    # Save results
    adata_results_path = OUTPUT_DIR / f"{out_prefix}_adata_all_methods.h5ad"
    adata.write(adata_results_path)

    # Save consensus scores as CSV
    consensus_scores_path = OUTPUT_DIR / f"{out_prefix}_consensus_scores.csv"
    scores_consensus.to_df().to_csv(consensus_scores_path)

    return {
        "message": f"Multiple enrichment methods completed with consensus scoring on {adata.n_obs} cells",
        "reference": "https://github.com/scverse/decoupler-tutorials/blob/main/example.ipynb",
        "artifacts": [
            {
                "description": "AnnData with all method scores",
                "path": str(adata_results_path.resolve()),
            },
            {
                "description": "Consensus enrichment scores",
                "path": str(consensus_scores_path.resolve()),
            },
            {
                "description": "Consensus scores heatmap",
                "path": str(consensus_heatmap_path.resolve()),
            },
        ],
    }


@example_mcp.tool
def decoupler_access_prior_knowledge(
    resource_name: Annotated[
        str | None, "Name of OmniPath resource to query (e.g., 'SIGNOR')"
    ] = None,
    out_prefix: Annotated[str | None, "Output file prefix"] = None,
) -> dict:
    """
    Query OmniPath resources to access prior knowledge networks for enrichment analysis.
    Input is resource name and output is available resources list and specific resource data if requested.
    """

    # Set output prefix
    if out_prefix is None:
        out_prefix = f"prior_knowledge_{timestamp}"

    # Get available resources
    resources = dc.op.show_resources()

    # Save resources list
    resources_path = OUTPUT_DIR / f"{out_prefix}_available_resources.csv"
    resources.to_csv(resources_path, index=False)

    artifacts = [
        {
            "description": "Available OmniPath resources",
            "path": str(resources_path.resolve()),
        }
    ]

    message = f"Found {len(resources)} available OmniPath resources"

    # If specific resource requested, download it
    if resource_name is not None:
        try:
            resource_data = dc.op.resource(resource_name)
            resource_data_path = OUTPUT_DIR / f"{out_prefix}_{resource_name}_network.csv"
            resource_data.to_csv(resource_data_path, index=False)

            artifacts.append(
                {
                    "description": f"{resource_name} network data",
                    "path": str(resource_data_path.resolve()),
                }
            )

            message += f" and downloaded {resource_name} with {len(resource_data)} interactions"

        except Exception as e:
            message += f" but failed to download {resource_name}: {str(e)}"

    return {
        "message": message,
        "reference": "https://github.com/scverse/decoupler-tutorials/blob/main/example.ipynb",
        "artifacts": artifacts,
    }