from scanpy.neighbors import Neighbors
from anndata import AnnData
from typing import Union, Optional, Any, Mapping, Callable, NamedTuple, Generator, Tuple
from numpy.random import RandomState
import numpy as np

Metric = Callable[[np.ndarray, np.ndarray], float]
import scanpy.logging as logg
N_DCS = 15  # default number of diffusion components

def neighbors(
    adata: AnnData,
    n_neighbors: int = 15,
    n_pcs: Optional[int] = None,
    use_rep: Optional[str] = None,
    knn: bool = True,
    mutual: bool = False,
    random_state: Optional[Union[int, RandomState]] = 0,
    method: str = 'umap',
    metric: Union[str, Metric] = 'euclidean',
    metric_kwds: Mapping[str, Any] = {},
    copy: bool = False,
) -> Optional[AnnData]:
    start = logg.info('computing neighbors')
    adata = adata.copy() if copy else adata
    if adata.isview:  # we shouldn't need this here...
        adata._init_as_actual(adata.copy())
    neighbors = Neighbors(adata)
    neighbors.compute_neighbors(
        n_neighbors=n_neighbors, knn=knn, n_pcs=n_pcs, use_rep=use_rep,
        method=method, metric=metric, metric_kwds=metric_kwds,
        random_state=random_state, write_knn_indices=mutual,
    )
    adata.uns['neighbors'] = {}
    adata.uns['neighbors']['params'] = {'n_neighbors': n_neighbors, 'method': method}
    adata.uns['neighbors']['params']['metric'] = metric
    if metric_kwds:
        adata.uns['neighbors']['params']['metric_kwds'] = metric_kwds
    if use_rep is not None:
        adata.uns['neighbors']['params']['use_rep'] = use_rep
    if n_pcs is not None:
        adata.uns['neighbors']['params']['n_pcs'] = n_pcs
    adata.uns['neighbors']['distances'] = neighbors.distances
    adata.uns['neighbors']['connectivities'] = neighbors.connectivities
    if mutual:
        import scipy.sparse as sp

        lm = sp.lil_matrix((adata.n_obs, adata.n_obs))
        lm.rows = neighbors.knn_indices
        lm.data = np.ones_like(neighbors.knn_indices)
        lm = lm.tocsr()
        lm.setdiag(0)
        lm.eliminate_zeros()
        lm = lm.multiply(lm.T) # build mnn mask

        adata.uns['neighbors']['distances'] = adata.uns['neighbors']['distances'].multiply(lm)
        adata.uns['neighbors']['connectivities'] = adata.uns['neighbors']['connectivities'].multiply(lm)

    if neighbors.rp_forest is not None:
        adata.uns['neighbors']['rp_forest'] = neighbors.rp_forest
    logg.info(
        '    finished',
        time=start,
        deep=(
            'added to `.uns[\'neighbors\']`\n'
            '    \'distances\', distances for each pair of neighbors\n'
            '    \'connectivities\', weighted adjacency matrix'
        ),
    )
    return adata if copy else None