# Biology imports.

import scanpy as sc

import mudata as md

# Typing imports.

from typing import Iterable

# Matrix operations.

import numpy as np

def top_features(

    mdata: md.MuData,

    mod: str = "rna",

    uns: str = "H_OT",

    dim: int = 0,

    threshold: float = 0.2,

) -> Iterable:

    """Returns the top features for a given modality and latent dimension.

    Args:

        mdata (md.MuData): The input data

        mod (str, optional): The modality. Defaults to 'rna'.

        uns (str, optional): Where to look for H. Defaults to 'H_OT'.

        dim (int, optional): The latent dimension. Defaults to 0.

        n_features (int, optional): The number of top features. Defaults to 5.

    Returns:

        Iterable: A list of features names.

    """

    # TODO: put variable names in uns!

    # Compue the number of features needed.

    n_features = np.sum(np.cumsum(np.sort(mdata[mod].uns[uns][:, dim])) > threshold)

    # Get names for highly variable features.

    idx = mdata[mod].var.highly_variable

    var_names = mdata[mod].var_names[idx]

    # Sort them by contribution.

    var_idx = np.argsort(mdata[mod].uns[uns][:, dim])[::-1]

    # Return the top ones.

    return var_names[var_idx[:n_features]].tolist()

def enrich(

    mdata: md.MuData,

    mod: str = "rna",

    uns: str = "H_OT",

    n_genes: int = 200,

    sources: Iterable[str] = ["GO:MF", "GO:CC", "GO:BP"],

    ordered: bool = True,

    domain_scope="custom_annotated",

):

    """Return Gene Set Enrichment Analysis results for each dimension.

    Args:

        mdata (md.MuData): Input data.

        mod (str, optional): Modality that contains genes. Defaults to 'rna'.

        uns (str, optional): Name of H matrix. Defaults to 'H_OT'.

        n_genes (int, optional):

            Number of top genes by dimension. Defaults to 200.

        sources (Iterable[str], optional):

            Enrichment sources. Defaults to ['GO:MF', 'GO:CC', 'GO:BP'].

        ordered (bool, optional):

            Make query with ordered genes. Defaults to True.

    Returns:

        Pandas dataframe with the results of the queries, as well

        as average best p_value across dimensions.

    """

    # Initialize ordered genes dictionary.

    ordered_genes = {}

    # Get background genes.

    background = mdata[mod].var.index.tolist()

    # For each dimension,

    for dim in range(mdata[mod].uns[uns].shape[1]):

        # Sort the gene indices by weight.

        idx_sorted = mdata[mod].uns[uns][:, dim].argsort()[::-1]

        if n_genes == "auto":

            nn = np.sum(np.cumsum(np.sort(mdata[mod].uns[uns][:, dim])) > 0.1)

        else:

            nn = n_genes

        # Select the `n_genes` highest genes.

        gene_list = mdata[mod].var[mdata[mod].var.highly_variable].index

        gene_list = gene_list[idx_sorted].tolist()[:nn]

        # Input them in the dictionary.

        ordered_genes["dimension " + str(dim)] = gene_list

    # Make the queries to gProfiler, specifying if genes are ordered.

    if "custom" in domain_scope:

        enr = sc.queries.enrich(

            ordered_genes,

            gprofiler_kwargs={

                "ordered": ordered,

                "sources": sources,

                "domain_scope": domain_scope,

                "background": background,

                "no_evidences": True,

            },

        )

    else:

        enr = sc.queries.enrich(

            ordered_genes,

            gprofiler_kwargs={

                "ordered": ordered,

                "sources": sources,

                "domain_scope": domain_scope,

                "no_evidences": True,

            },

        )

    # Compute the average of the best p_values for each dimension.

    mean_best_p = enr.groupby("query")["p_value"].min().mean()

    # Return the results of the queries and the average best p_value.

    return enr, mean_best_p