import anndata as ad
import mudata as md
import numpy as np
import pandas as pd
import scanpy as sc
import seaborn as sns
from matplotlib import pyplot as plt


def clustermap(mdata: md.MuData, obsm: str = "W_OT", cmap="viridis", **kwds):
    """Wrapper around Scanpy's clustermap.

    Args:
        mdata (md.MuData): The input data
        obsm (str, optional): The obsm field to consider. Defaults to 'W_OT'.
        cmap (str, optional): The colormap. Defaults to 'viridis'.
    """

    # Create an AnnData with the joint embedding.
    joint_embedding = ad.AnnData(mdata.obsm[obsm], obs=mdata.obs)

    # Make the clustermap plot.
    sc.pl.clustermap(joint_embedding, cmap=cmap, **kwds)


def factor_violin(
    mdata: md.MuData,
    groupby: str,
    obsm: str = "W_OT",
    dim: int = 0,
    **kwds,
):
    """Make a violin plot of cells for a given latent dimension.

    Args:
        mdata (md.MuData): The input data
        dim (int, optional): The latent dimension. Defaults to 0.
        obsm (str, optional): The embedding. Defaults to 'W_OT'.
        groupby (str, optional): Observation groups.
    """

    # Create an AnnData with the joint embedding.
    joint_embedding = ad.AnnData(mdata.obsm[obsm], obs=mdata.obs)

    # Add the obs field that we're interested in.
    joint_embedding.obs["Factor " + str(dim)] = joint_embedding.X[:, dim]

    # Make the violin plot.
    sc.pl.violin(joint_embedding, keys="Factor " + str(dim), groupby=groupby, **kwds)


def heatmap(
    mdata: md.MuData,
    groupby: str,
    obsm: str = "W_OT",
    cmap: str = "viridis",
    sort_var: bool = False,
    save: str = None,
    **kwds,
) -> None:
    """Produce a heatmap of an embedding

    Args:
        mdata (md.MuData): Input data
        groupby (str): What to group by
        obsm (str): The embedding. Defaults to 'W_OT'.
        cmap (str, optional): Color map. Defaults to 'viridis'.
        sort_var (bool, optional):
            Sort dimensions by variance. Defaults to False.
    """

    # Create an AnnData with the joint embedding.
    joint_embedding = ad.AnnData(mdata.obsm[obsm], obs=mdata.obs)

    # Try to compute a dendrogram.
    try:
        sc.pp.pca(joint_embedding)
        sc.tl.dendrogram(joint_embedding, groupby=groupby, use_rep="X_pca")
    except Exception:
        print("Dendrogram not computed.")
        pass

    # Get the dimension names to show.
    if sort_var:
        idx = joint_embedding.X.std(0).argsort()[::-1]
        var_names = joint_embedding.var_names[idx]
    else:
        var_names = joint_embedding.var_names

    # PLot the heatmap.
    return sc.pl.heatmap(
        joint_embedding, var_names, groupby=groupby, cmap=cmap, save=save, **kwds
    )


def enrich(enr: pd.DataFrame, query_name: str, n_terms: int = 10):
    """Display a list of enriched terms.

    Args:
        enr (pd.DataFrame): The enrichment object returned by mowgli.tl.enrich()
        query_name (str): The name of the query, e.g. "dimension 0".
    """

    # Subset the enrichment object to the query of interest.
    sub_enr = enr[enr["query"] == query_name].head(n_terms)
    sub_enr["minlogp"] = -np.log10(sub_enr["p_value"])

    fig, ax = plt.subplots()

    # Display the enriched terms.
    ax.hlines(
        y=sub_enr["name"],
        xmin=0,
        xmax=sub_enr["minlogp"],
        color="lightgray",
        zorder=1,
        alpha=0.8,
    )
    sns.scatterplot(
        data=sub_enr,
        x="minlogp",
        y="name",
        hue="source",
        s=100,
        alpha=0.8,
        ax=ax,
        zorder=3,
    )

    ax.set_xlabel("$-log_{10}(p)$")
    ax.set_ylabel("Enriched terms")

    plt.show()


def top_features(
    mdata: md.MuData,
    mod: str = "rna",
    uns: str = "H_OT",
    dim: int = 0,
    n_top: int = 10,
    ax: plt.axes = None,
    palette: str = "Blues_r",
):
    """Display the top features for a given dimension.

    Args:
        mdata (md.MuData): The input mdata object
        mod (str, optional): The modality to consider. Defaults to 'rna'.
        uns (str, optional): The uns field to consider. Defaults to 'H_OT'.
        dim (int, optional): The latent dimension. Defaults to 0.
        n_top (int, optional): The number of top features to display. Defaults to 10.
        ax (plt.axes, optional): The axes to use. Defaults to None.
        palette (str, optional): The color palette to use. Defaults to 'Blues_r'.

    Returns:
        plt.axes: The axes used.
    """

    # Get the variable names.
    var_names = mdata[mod].var_names[mdata[mod].var.highly_variable]

    # Get the top features.
    idx_top_features = np.argsort(mdata[mod].uns[uns][:, dim])[::-1][:n_top]
    df = pd.DataFrame(
        {
            "features": var_names[idx_top_features],
            "weights": mdata[mod].uns[uns][idx_top_features, dim],
        }
    )

    # Display the top features.
    if ax is None:
        ax = sns.barplot(data=df, x="weights", y="features", palette=palette)
    else:
        sns.barplot(data=df, x="weights", y="features", palette=palette, ax=ax)

    return ax


def umap(
    mdata: md.MuData,
    dim: int | list = 0,
    rescale: bool = False,
    obsm: str = "W_OT",
    neighbours_key=None,
    **kwds,
):
    """Wrapper around Scanpy's sc.pl.umap. Computes UMAP for a given latent dimension and plots it.
    Args:
        mdata (md.MuData): The input data
        dim (int | list, optional): The latent dimension. Defaults to 0.
        rescale (bool, optional): If True, Rescale the color palette across all plots to the maximum value in the weight matrix. Defaults to False.
        obsm (str, optional): The embedding. Defaults to 'W_OT'.
        neighbours_key (str, optional): The key for the neighbours in `mdata.uns` to use to compute neighbors. Defaults to None.
    """

    adata_tmp = ad.AnnData(mdata.obsm[obsm], obs=pd.DataFrame(index=mdata.obs.index))

    if isinstance(dim, int):
        mowgli_cat = f"mowgli:{dim}"

    elif isinstance(dim, list):
        # clean dim of doubles and sort them
        dim = sorted(list(set(dim)))
        mowgli_cat = [f"mowgli:{x}" for x in dim]

    else:
        raise ValueError("dim must be an integer or a list of integers")

    adata_tmp.obs[mowgli_cat] = adata_tmp.X[:, dim]

    # check if neighbors exists
    if neighbours_key is None:
        print("Computing neighbors with scanpy default parameters")
        neighbours_key = "mowgli_neighbors"  # set the default neighbors key
        # compute neiughborts using all dimension in the mowgli matrix
        sc.pp.neighbors(adata_tmp, use_rep="X", key_added=neighbours_key)

    else:
        if neighbours_key not in mdata.uns.keys():
            raise ValueError(f"neighbours key {neighbours_key} not found in mdata.uns")

        adata_tmp.uns[neighbours_key] = mdata.uns[neighbours_key]

    # compute umap
    print("Computing UMAP")
    sc.tl.umap(adata_tmp, neighbors_key=neighbours_key)

    # plot umap
    if rescale:
        vmax = adata_tmp.X.max()
        sc.pl.umap(adata_tmp, color=mowgli_cat, size=18.5, alpha=0.4, vmax=vmax, **kwds)
    else:
        sc.pl.umap(adata_tmp, color=mowgli_cat, size=18.5, alpha=0.4, **kwds)