from __future__ import annotations

from typing import Literal
import scipy.sparse as ss
import numpy as np

import snapatac2
import snapatac2._snapatac2 as internal
from snapatac2._utils import get_igraph_from_adjacency, is_anndata

def leiden(
    adata: internal.AnnData | internal.AnnDataSet | ss.spmatrix,
    resolution: float = 1,
    objective_function: Literal["CPM", "modularity"] = "modularity",
    min_cluster_size: int = 5,
    n_iterations: int = -1,
    random_state: int = 0,
    key_added: str = "leiden",
    weighted: bool = False,
    inplace: bool = True,
) -> np.ndarray | None:
    """
    Cluster cells into subgroups [Traag18]_.

    Cluster cells using the Leiden algorithm [Traag18]_,
    an improved version of the Louvain algorithm [Blondel08]_.
    It has been proposed for single-cell analysis by [Levine15]_.
    This requires having ran :func:`~snapatac2.pp.knn`.

    Parameters
    ----------
    adata
        The annotated data matrix or sparse adjacency matrix of the graph,
        defaults to neighbors connectivities.
    resolution
        A parameter value controlling the coarseness of the clustering.
        Higher values lead to more clusters.
        Set to `None` if overriding `partition_type`
        to one that doesn't accept a `resolution_parameter`.
    objective_function
        whether to use the Constant Potts Model (CPM) or modularity.
        Must be either "CPM" or "modularity".
    min_cluster_size
        The minimum size of clusters.
    n_iterations
        How many iterations of the Leiden clustering algorithm to perform.
        Positive values above 2 define the total number of iterations to perform,
        -1 has the algorithm run until it reaches its optimal clustering.
    random_state
        Change the initialization of the optimization.
    key_added
        `adata.obs` key under which to add the cluster labels.
    weighted
        Whether to use the edge weights in the graph
    inplace
        Whether to store the result in the anndata object.

    Returns
    -------
    np.ndarray | None
        If `inplace=True`, update `adata.obs[key_added]` to store an array of
        dim (number of samples) that stores the subgroup id
        (`'0'`, `'1'`, ...) for each cell. Otherwise, returns the array directly.
    """
    from igraph import set_random_number_generator
    from collections import Counter
    import polars
    import random

    random.seed(random_state)
    set_random_number_generator(random)

    if is_anndata(adata):
        adjacency = adata.obsp["distances"]
    else:
        inplace = False
        adjacency = adata

    gr = get_igraph_from_adjacency(adjacency)

    if weighted:
        weights = np.array(gr.es["weight"])
        weights = np.exp(-weights)
    else:
        weights = None

    groups = gr.community_leiden(
        objective_function=objective_function,
        weights=weights,
        resolution=resolution,
        beta=0.01,
        initial_membership=None,
        n_iterations=n_iterations,
    ).membership

    new_cl_id = dict(
        [
            (cl, i) if count >= min_cluster_size else (cl, -1)
            for (i, (cl, count)) in enumerate(Counter(groups).most_common())
        ]
    )
    for i in range(len(groups)):
        groups[i] = new_cl_id[groups[i]]

    groups = np.array(groups, dtype=np.str_)
    if inplace:
        adata.obs[key_added] = polars.Series(
            groups,
            dtype=polars.datatypes.Categorical,
        )
    else:
        return groups


def leiden_sweep(
    adata: internal.AnnData | internal.AnnDataSet | ss.spmatrix,
    resolutions: list[float],
    use_rep: str | np.ndarray = "X_spectral",
    objective_function: Literal["CPM", "modularity", "RBConfiguration"] = "modularity",
    min_cluster_size: int = 5,
    n_iterations: int = -1,
    random_state: int = 0,
    weighted: bool = False,
    n_jobs: int = 16,
):
    """
    Perform a sweep over multiple resolutions for Leiden clustering and compute silhouette scores.

    Parameters
    ----------
    adata
        The annotated data matrix or sparse adjacency matrix of the graph,
        defaults to neighbors connectivities.
    resolutions
        A list of resolution values to evaluate.
    use_rep
        Which data in `adata.obsm` to use for clustering. Default is "X_spectral".
    objective_function
        whether to use the Constant Potts Model (CPM) or modularity.
        Must be either "CPM", "modularity" or "RBConfiguration".
    min_cluster_size
        The minimum size of clusters.
    n_iterations
        How many iterations of the Leiden clustering algorithm to perform.
        Positive values above 2 define the total number of iterations to perform,
        -1 has the algorithm run until it reaches its optimal clustering.
    random_state
        Change the initialization of the optimization.
    weighted
        Whether to use the edge weights in the graph
    n_jobs
        The number of parallel jobs to run.

    Returns
    -------
    list[dict]
        A dictionary mapping each resolution to its corresponding silhouette score.
    """
    from sklearn.metrics import silhouette_score
    from multiprocess import get_context

    if is_anndata(adata):
        distances = adata.obsp["distances"]
        mat = adata.obsm[use_rep]
    else:
        mat = use_rep
        distances = adata

    def _func(resolution):
        groups = leiden(
            distances,
            resolution,
            objective_function=objective_function,
            min_cluster_size=min_cluster_size,
            n_iterations=n_iterations,
            random_state=random_state,
            weighted=weighted,
            inplace=False,
        )
        if len(set(groups)) > 1:
            score = silhouette_score(
                mat,
                groups,
                sample_size=20000,
            )
        else:
            score = 0
        return  {
            "resolution": resolution,
            "n_clusters": len(set(groups)),
            "silhouette_score": score,
        }

    with get_context("spawn").Pool(n_jobs) as p:
        return list(p.imap(_func, resolutions))

def kmeans(
    adata: internal.AnnData | internal.AnnDataSet | np.ndarray,
    n_clusters: int,
    n_iterations: int = -1,
    random_state: int = 0,
    use_rep: str = "X_spectral",
    key_added: str = "kmeans",
    inplace: bool = True,
) -> np.ndarray | None:
    """
    Cluster cells into subgroups using the K-means algorithm, a classical algorithm in data mining.


    Parameters
    ----------
    adata
        The annotated data matrix.
    n_clusters
        Number of clusters to return.
    n_iterations
        How many iterations of the kmeans clustering algorithm to perform.
        Positive values above 2 define the total number of iterations to perform,
        -1 has the algorithm run until it reaches its optimal clustering.
    random_state
        Change the initialization of the optimization.
    use_rep
        Which data in `adata.obsm` to use for clustering. Default is "X_spectral".
    key_added
        `adata.obs` key under which to add the cluster labels.

    Returns
    -------
    adds fields to `adata`:
    `adata.obs[key_added]`
        Array of dim (number of samples) that stores the subgroup id
        (`'0'`, `'1'`, ...) for each cell.
    `adata.uns['kmeans']['params']`
        A dict with the values for the parameters `n_clusters`, `random_state`,
        and `n_iterations`.
    """
    import polars

    if is_anndata(adata):
        data = adata.obsm[use_rep]
    else:
        data = adata
    groups = internal.kmeans(n_clusters, data)
    groups = np.array(groups, dtype=np.str_)
    if inplace:
        adata.obs[key_added] = polars.Series(
            groups,
            dtype=polars.datatypes.Categorical,
        )
        # store information on the clustering parameters
        # adata.uns['kmeans'] = {}
        # adata.uns['kmeans']['params'] = dict(
        #    n_clusters=n_clusters,
        #    random_state=random_state,
        #    n_iterations=n_iterations,
        # )

    else:
        return groups


def hdbscan(
    adata: internal.AnnData,
    min_cluster_size: int = 5,
    min_samples: int | None = None,
    cluster_selection_epsilon: float = 0.0,
    alpha: float = 1.0,
    cluster_selection_method: str = "eom",
    random_state: int = 0,
    use_rep: str = "X_spectral",
    key_added: str = "hdbscan",
    **kwargs,
) -> None:
    """
    Cluster cells into subgroups using the HDBSCAN algorithm.

    Parameters
    ----------
    adata
        The annotated data matrix.
    min_cluster_size
        The minimum size of clusters; single linkage splits that contain
        fewer points than this will be considered points "falling out" of
        a cluster rather than a cluster splitting into two new clusters.
    min_samples
        The number of samples in a neighbourhood for a point to be considered a core point.
    cluster_selection_epsilon
        A distance threshold. Clusters below this value will be merged.
    alpha
        A distance scaling parameter as used in robust single linkage.
    cluster_selection_method
        The method used to select clusters from the condensed tree.
        The standard approach for HDBSCAN* is to use an Excess of Mass
        algorithm to find the most persistent clusters.
        Alternatively you can instead select the clusters at the leaves of
        the tree - this provides the most fine grained and homogeneous clusters.
        Options are: "eom" or "leaf".
    random_state
        Change the initialization of the optimization.
    use_rep
        Which data in `adata.obsm` to use for clustering. Default is "X_spectral".
    key_added
        `adata.obs` key under which to add the cluster labels.

    Returns
    -------
    adds fields to `adata`:
    `adata.obs[key_added]`
        Array of dim (number of samples) that stores the subgroup id
        (`'0'`, `'1'`, ...) for each cell.
    """
    import pandas as pd
    import hdbscan

    data = adata.obsm[use_rep][...]
    clusterer = hdbscan.HDBSCAN(
        min_cluster_size=min_cluster_size,
        min_samples=min_samples,
        cluster_selection_epsilon=cluster_selection_epsilon,
        alpha=alpha,
        cluster_selection_method=cluster_selection_method,
        **kwargs,
    )
    clusterer.fit(data)
    groups = clusterer.labels_
    adata.obs[key_added] = pd.Categorical(
        values=groups.astype("U"),
        categories=sorted(map(str, np.unique(groups))),
    )


def dbscan(
    adata: internal.AnnData,
    eps: float = 0.5,
    min_samples: int = 5,
    leaf_size: int = 30,
    n_jobs: int | None = None,
    use_rep: str = "X_spectral",
    key_added: str = "dbscan",
) -> None:
    """
    Cluster cells into subgroups using the DBSCAN algorithm.

    Parameters
    ----------
    adata
        The annotated data matrix.
    eps
        The maximum distance between two samples for one to be considered as
        in the neighborhood of the other. This is not a maximum bound on the
        distances of points within a cluster. This is the most important
        DBSCAN parameter to choose appropriately for your data set and distance function.
    min_samples
        The number of samples (or total weight) in a neighborhood for a point
        to be considered as a core point. This includes the point itself.
    leaf_size
        Leaf size passed to BallTree or cKDTree. This can affect the speed of
        the construction and query, as well as the memory required to store the
        tree. The optimal value depends on the nature of the problem.
    n_jobs
        The number of parallel jobs to run. None means 1 unless in a
        joblib.parallel_backend context. -1 means using all processors.
    use_rep
        Which data in `adata.obsm` to use for clustering. Default is "X_spectral".
    key_added
        `adata.obs` key under which to add the cluster labels.

    Returns
    -------
    adds fields to `adata`:
    `adata.obs[key_added]`
        Array of dim (number of samples) that stores the subgroup id
        (`'0'`, `'1'`, ...) for each cell.
    """
    from sklearn.cluster import DBSCAN
    import pandas as pd

    data = adata.obsm[use_rep][...]

    clustering = DBSCAN(
        eps=eps,
        min_samples=min_samples,
        metric="euclidean",
        leaf_size=leaf_size,
        n_jobs=n_jobs,
    ).fit(data)
    groups = clustering.labels_
    adata.obs[key_added] = pd.Categorical(
        values=groups.astype("U"),
        categories=sorted(map(str, np.unique(groups))),
    )
