#!/usr/bin/env python3
# coding: utf-8
"""
@author: Yiran Wu  wuyiran@genomics.cn
@last modified by: Yiran Wu
@file:neighbors.py
@time:2021/09/01
"""
from types import MappingProxyType
from typing import (
    Union,
    Any,
    Mapping,
    Optional,
    Callable
)

import igraph as ig
import numpy as np
from numpy import random
from scipy.sparse import (
    issparse,
    coo_matrix,
    csr_matrix
)
from sklearn.metrics import pairwise_distances
from sklearn.neighbors import NearestNeighbors
from sklearn.utils import check_random_state
from typing_extensions import Literal

from ..log_manager import logger

AnyRandom = Union[None, int, random.RandomState]

_Method = Literal['umap', 'gauss', 'rapids']
_MetricFn = Callable[[np.ndarray, np.ndarray], float]
# from sklearn.metrics.pairwise_distances.__doc__:
_MetricSparseCapable = Literal[
    'cityblock', 'cosine', 'euclidean', 'l1', 'l2', 'manhattan'
]
_MetricScipySpatial = Literal[
    'braycurtis',
    'canberra',
    'chebyshev',
    'correlation',
    'dice',
    'hamming',
    'jaccard',
    'kulsinski',
    'mahalanobis',
    'minkowski',
    'rogerstanimoto',
    'russellrao',
    'seuclidean',
    'sokalmichener',
    'sokalsneath',
    'sqeuclidean',
    'yule',
]
_Metric = Union[_MetricSparseCapable, _MetricScipySpatial]


def compute_neighbors_rapids(
        X: np.ndarray, n_neighbors: int, metric: _Metric = 'euclidean'
):
    """Compute nearest neighbors using RAPIDS cuml.

    Parameters
    ----------
    X: array of shape (n_samples, n_features)
        The data to compute nearest neighbors for.
    n_neighbors
        The number of neighbors to use.
    metric
        The metric to use to compute distances in high dimensional space.
        This string must match a valid predefined metric in RAPIDS cuml.

        Returns
    -------
    **knn_indices**, **knn_dists** : np.arrays of shape (n_observations, n_neighbors)
    """
    try:
        from cuml.neighbors.nearest_neighbors import NearestNeighbors
    except ImportError:
        raise ImportError("Your env don't have GPU related RAPIDS packages, if you want to run this option, follow the "
                          "guide at https://stereopy.readthedocs.io/en/latest/Tutorials/clustering_by_gpu.html")

    nn = NearestNeighbors(n_neighbors=n_neighbors, metric=metric)
    X_contiguous = np.ascontiguousarray(X, dtype=np.float32)
    nn.fit(X_contiguous)
    knn_dist, knn_indices = nn.kneighbors(X_contiguous)
    logger.info('cuml NearestNeighbors run end')
    return knn_indices, knn_dist


def find_neighbors(
        x: np.ndarray,
        n_neighbors: Optional[int] = 15,
        n_pcs: Optional[int] = None,
        method: Optional[_Method] = 'umap',
        metric: Union[_Metric, _MetricFn] = 'euclidean',
        metric_kwds: Mapping[str, Any] = MappingProxyType({}),
        knn: bool = True,
        random_state: AnyRandom = 0,
        n_jobs: int = 10
):
    """

    :param x:
        class:`~numpy.ndarray`. PCA representation of data.
    :param n_neighbors:
        Use this number of nearest neighbors.
    :param n_pcs:
    :param method:
        Use 'umap' or 'gauss'. for computing connectivities.
    :param metric:
        A known metric's name or a callable that returns a distance.
        include:
            * euclidean
            * manhattan
            * chebyshev
            * minkowski
            * canberra
            * braycurtis
            * mahalanobis
            * wminkowski
            * seuclidean
            * cosine
            * correlation
            * haversine
            * hamming
            * jaccard
            * dice
            * russelrao
            * kulsinski
            * rogerstanimoto
            * sokalmichener
            * sokalsneath
            * yule
    :param metric_kwds:
        Arguments to pass on to the metric, such as the ``p`` value for Minkowski distance.
    :param knn:
        If `True`, use a hard threshold to restrict the number of neighbors to
        `n_neighbors`, that is, consider a knn graph. Otherwise, use a Gaussian
        Kernel to assign low weights to neighbors more distant than the
        `n_neighbors` nearest neighbor.
    :param random_state:
        A state capable being used as a numpy random state.
    :return:
        neighbor: Neighbors object
        dists: sparse
        connectivities: sparse
    """
    neighbor = Neighbors(
        x=x,
        n_neighbors=n_neighbors,
        n_pcs=n_pcs,
        metric=metric,
        method=method,
        knn=knn,
        random_state=random_state,
    )
    neighbor.check_setting()
    neighbor.x = neighbor.choose_x()
    use_dense_distances = (neighbor.metric == 'euclidean' and neighbor.x.shape[0] < 8192) or not neighbor.knn
    dists = neighbor.x

    connectivities = None
    if use_dense_distances:
        dists = pairwise_distances(neighbor.x, metric=neighbor.metric, **metric_kwds)
        knn_indices, knn_distances = neighbor.get_indices_distances_from_dense_matrix(dists)
        if knn:
            dists = neighbor.get_parse_distances_numpy(
                knn_indices, knn_distances, neighbor.x.shape[0],
            )
    elif method == 'rapids':
        knn_indices, knn_distances = compute_neighbors_rapids(
            neighbor.x, n_neighbors, metric=metric
        )
    else:
        if neighbor.x.shape[0] < 4096:
            dists = pairwise_distances(neighbor.x, metric=neighbor.metric, **metric_kwds)
            neighbor.metric = 'precomputed'
        knn_indices, knn_distances, forest = neighbor.compute_neighbors_umap(
            dists, random_state=neighbor.random_state, metric_kwds=metric_kwds, n_jobs=n_jobs)

    if not use_dense_distances or neighbor.method in {'umap', 'rapids'}:
        connectivities = neighbor.compute_connectivities_umap(knn_indices, knn_distances)
        dists = neighbor.get_parse_distances_umap(knn_indices, knn_distances, )

    if method == 'gauss':
        connectivities = neighbor.compute_connectivities_diffmap(dists)

    return neighbor, dists, connectivities


class Neighbors(object):
    def __init__(self, x, n_neighbors, n_pcs, method, metric, knn, random_state):
        self.x = x
        self.n_neighbors = n_neighbors
        self.n_pcs = n_pcs
        self.metric = metric
        self.method = method
        self.knn = knn
        self.random_state = random_state

    def check_setting(self):
        if self.method == 'umap' and not self.knn:
            logger.error('`method=\'umap\' only with knn=True`.')
        if self.method not in {'umap', 'gauss', 'rapids'}:
            logger.error('`method=\'umap\' or method=\'gauss\'`.')

    def choose_x(self):
        self.x = self.x[:, :self.n_pcs]
        return self.x

    def get_indices_distances_from_dense_matrix(self, dists, ):
        sample_range = np.arange(dists.shape[0])[:, None]
        indices = np.argpartition(dists, self.n_neighbors - 1, axis=1)[:, :self.n_neighbors]
        indices = indices[sample_range, np.argsort(dists[sample_range, indices])]
        distances = dists[sample_range, indices]
        return indices, distances

    def get_indices_distances_from_sparse_matrix(self, dists):
        indices = np.zeros((dists.shape[0], self.n_neighbors), dtype=int)
        distances = np.zeros((dists.shape[0], self.n_neighbors), dtype=dists.dtype)
        n_neighbors_m1 = self.n_neighbors - 1
        for i in range(indices.shape[0]):
            neighbors = dists[i].nonzero()  # 'true' and 'spurious' zeros
            indices[i, 0] = i
            distances[i, 0] = 0
            # account for the fact that there might be more than n_neighbors
            # due to an approximate search
            # [the point itself was not detected as its own neighbor during the search]
            if len(neighbors[1]) > n_neighbors_m1:
                sorted_indices = np.argsort(dists[i][neighbors].A1)[:n_neighbors_m1]
                indices[i, 1:] = neighbors[1][sorted_indices]
                distances[i, 1:] = dists[i][
                    neighbors[0][sorted_indices], neighbors[1][sorted_indices]
                ]
            else:
                indices[i, 1:] = neighbors[1]
                distances[i, 1:] = dists[i][neighbors]
        return indices, distances

    def get_parse_distances_numpy(self, indices, distances, n_obs, ):
        n_nonzero = n_obs * self.n_neighbors
        indptr = np.arange(0, n_nonzero + 1, self.n_neighbors)
        dists = csr_matrix(
            (
                distances.copy().ravel(),  # copy the data, otherwise strange behavior here
                indices.copy().ravel(),
                indptr,
            ),
            shape=(n_obs, n_obs),
        )
        dists.eliminate_zeros()
        return dists

    def get_parse_distances_umap(self, nn_idx, nn_dist):
        n_obs = self.x.shape[0]
        rows = np.zeros((n_obs * self.n_neighbors), dtype=np.int64)
        cols = np.zeros((n_obs * self.n_neighbors), dtype=np.int64)
        vals = np.zeros((n_obs * self.n_neighbors), dtype=np.float64)

        for i in range(nn_idx.shape[0]):
            for j in range(self.n_neighbors):
                if nn_idx[i, j] == -1:
                    continue  # We didn't get the full knn for i
                if nn_idx[i, j] == i:
                    val = 0.0
                else:
                    val = nn_dist[i, j]

                rows[i * self.n_neighbors + j] = i
                cols[i * self.n_neighbors + j] = nn_idx[i, j]
                vals[i * self.n_neighbors + j] = val

        distances = coo_matrix((vals, (rows, cols)), shape=(n_obs, n_obs))
        distances.eliminate_zeros()
        return distances.tocsr()

    def compute_neighbors_umap(
            self,
            x,
            random_state: AnyRandom = None,
            angular: bool = False,
            verbose: bool = False,
            metric_kwds: Mapping[str, Any] = MappingProxyType({}),
            n_jobs: int = 10
    ):
        from umap.umap_ import nearest_neighbors
        random_state = check_random_state(random_state)
        knn_indices, knn_dists, forest = nearest_neighbors(
            x,
            self.n_neighbors,
            random_state=random_state,
            metric=self.metric,
            metric_kwds=metric_kwds,
            angular=angular,
            verbose=verbose,
            n_jobs=n_jobs
        )

        return knn_indices, knn_dists, forest

    def find_n_neighbors(self, ):
        nbrs = NearestNeighbors(n_neighbors=self.n_neighbors + 1, algorithm='ball_tree').fit(self.x)
        dists, indices = nbrs.kneighbors(self.x)
        nn_idx = indices[:, 1:]
        nn_dist = dists[:, 1:]
        return nn_idx, nn_dist

    def get_igraph_from_knn(self, nn_idx, nn_dist):
        j = nn_idx.ravel().astype(int)
        dist = nn_dist.ravel()
        i = np.repeat(np.arange(nn_idx.shape[0]), self.n_neighbors)

        vertex = list(range(nn_dist.shape[0]))
        edges = list(tuple(zip(i, j)))
        g = ig.Graph()
        g.add_vertices(vertex)
        g.add_edges(edges)
        g.es['weight'] = dist
        return g

    @staticmethod
    def get_igraph_from_adjacency(adjacency, directed=None):
        """Get igraph graph from adjacency matrix."""
        sources, targets = adjacency.nonzero()
        weights = adjacency[sources, targets]
        if isinstance(weights, np.matrix):
            weights = weights.A1
        g = ig.Graph(directed=directed)
        g.add_vertices(adjacency.shape[0])  # this adds adjacency.shape[0] vertices
        g.add_edges(list(zip(sources, targets)))
        try:
            g.es['weight'] = weights
        except KeyError:
            pass
        if g.vcount() != adjacency.shape[0]:
            logger.error(
                f'The constructed graph has only {g.vcount()} nodes. '
                'Your adjacency matrix contained redundant nodes.'
            )
        return g

    def compute_connectivities_umap(self, nn_idx, nn_dist):
        from umap.umap_ import fuzzy_simplicial_set
        n_obs = self.x.shape[0]
        x = coo_matrix(([], ([], [])), shape=(n_obs, 1))
        connectivities = fuzzy_simplicial_set(x, self.n_neighbors, None, None, knn_indices=nn_idx, knn_dists=nn_dist,
                                              set_op_mix_ratio=1.0, local_connectivity=1.0)
        if isinstance(connectivities, tuple):
            connectivities = connectivities[0]

        return connectivities.tocsr()

    def compute_connectivities_diffmap(self, dists, ):
        # init distances
        if self.knn:
            dsq = dists.power(2)
            indices, distances_sq = self.get_indices_distances_from_sparse_matrix(dsq, )
        else:
            dsq = np.power(dists, 2)
            indices, distances_sq = self.get_indices_distances_from_dense_matrix(dsq, )

        # exclude the first point, the 0th neighbor
        indices = indices[:, 1:]
        distances_sq = distances_sq[:, 1:]

        # choose sigma, the heuristic here doesn't seem to make much of a difference,
        # but is used to reproduce the figures of Haghverdi et al. (2016)
        if self.knn:
            # as the distances are not sorted
            # we have decay within the n_neighbors first neighbors
            sigmas_sq = np.median(distances_sq, axis=1)
        else:
            # the last item is already in its sorted position through argpartition
            # we have decay beyond the n_neighbors neighbors
            sigmas_sq = distances_sq[:, -1] / 4
        sigmas = np.sqrt(sigmas_sq)

        # compute the symmetric weight matrix
        if not issparse(dists):
            Num = 2 * np.multiply.outer(sigmas, sigmas)
            Den = np.add.outer(sigmas_sq, sigmas_sq)
            W = np.sqrt(Num / Den) * np.exp(-dsq / Den)
            # make the weight matrix sparse
            if not self.knn:
                mask = W > 1e-14
                W[~mask] = 0
            else:
                # restrict number of neighbors to ~k
                # build a symmetric mask
                mask = np.zeros(dsq.shape, dtype=bool)
                for i, row in enumerate(indices):
                    mask[i, row] = True
                    for j in row:
                        if i not in set(indices[j]):
                            W[j, i] = W[i, j]
                            mask[j, i] = True
                # set all entries that are not nearest neighbors to zero
                W[~mask] = 0
        else:
            W = (
                dsq.copy()
            )  # need to copy the distance matrix here; what follows is inplace
            for i in range(len(dsq.indptr[:-1])):
                row = dsq.indices[dsq.indptr[i]: dsq.indptr[i + 1]]
                num = 2 * sigmas[i] * sigmas[row]
                den = sigmas_sq[i] + sigmas_sq[row]
                W.data[dsq.indptr[i]: dsq.indptr[i + 1]] = np.sqrt(num / den) * np.exp(
                    -dsq.data[dsq.indptr[i]: dsq.indptr[i + 1]] / den
                )
            W = W.tolil()
            for i, row in enumerate(indices):
                for j in row:
                    if i not in set(indices[j]):
                        W[j, i] = W[i, j]
            W = W.tocsr()
        connectivities = W
        return connectivities
