# Path: umap/utils.py
# @numba.njit(parallel=True)
# def submatrix(dmat, indices_col, n_neighbors):
#     """Return a submatrix given an orginal matrix and the indices to keep.
# 
#     Parameters
#     ----------
#     dmat: array, shape (n_samples, n_samples)
#         Original matrix.
# 
#     indices_col: array, shape (n_samples, n_neighbors)
#         Indices to keep. Each row consists of the indices of the columns.
# 
#     n_neighbors: int
#         Number of neighbors.
# 
#     Returns
#     -------
#     submat: array, shape (n_samples, n_neighbors)
#         The corresponding submatrix.
#     """
#     n_samples_transform, n_samples_fit = dmat.shape
#     submat = np.zeros((n_samples_transform, n_neighbors), dtype=dmat.dtype)
#     for i in numba.prange(n_samples_transform):
#         for j in numba.prange(n_neighbors):
#             submat[i, j] = dmat[i, indices_col[i, j]]
#     return submat
# 
# def average_nn_distance(dist_matrix):
#     """Calculate the average distance to each points nearest neighbors.
# 
#     Parameters
#     ----------
#     dist_matrix: a csr_matrix
#         A distance matrix (usually umap_model.graph_)
# 
#     Returns
#     -------
#     An array with the average distance to each points nearest neighbors
# 
#     """
#     (row_idx, col_idx, val) = scipy.sparse.find(dist_matrix)
# 
#     # Count/sum is done per row
#     count_non_zero_elems = np.bincount(row_idx)
#     sum_non_zero_elems = np.bincount(row_idx, weights=val)
#     averages = sum_non_zero_elems / count_non_zero_elems
# 
#     if any(np.isnan(averages)):
#         warn(
#             "Embedding contains disconnected vertices which will be ignored."
#             "Use umap.utils.disconnected_vertices() to identify them."
#         )
# 
#     return averages


import numpy as np
import numba
    import pandas as pd
    import datashader as ds
    import datashader.transfer_functions as tf
    import datashader.bundling as bd
    import matplotlib.pyplot as plt
    import colorcet
    import matplotlib.colors
    import matplotlib.cm
    import bokeh.plotting as bpl
    import bokeh.transform as btr
    import holoviews as hv
    import holoviews.operation.datashader as hd
import sklearn.decomposition
import sklearn.cluster
import sklearn.neighbors
from warnings import warn
from matplotlib.patches import Patch
from umap.utils import submatrix, average_nn_distance
from bokeh.plotting import show as show_interactive
from bokeh.plotting import output_file, output_notebook
from bokeh.layouts import column
from bokeh.models import CustomJS, TextInput
from matplotlib.pyplot import show as show_static
from warnings import warn

def _get_metric(umap_object):
    if hasattr(umap_object, "metric"):
        return umap_object.metric
    else:
        # Assume euclidean if no attribute per cuML.UMAP
        return "euclidean"


def _get_metric_kwds(umap_object):
    if hasattr(umap_object, "_metric_kwds"):
        return umap_object._metric_kwds
    else:
        # Assume no keywords exist
        return {}


def _embed_datashader_in_an_axis(datashader_image, ax):
    img_rev = datashader_image.data[::-1]
    mpl_img = np.dstack([_blue(img_rev), _green(img_rev), _red(img_rev)])
    ax.imshow(mpl_img)
    return ax


def _nhood_search(umap_object, nhood_size):
    if hasattr(umap_object, "_small_data") and umap_object._small_data:
        dmat = sklearn.metrics.pairwise_distances(umap_object._raw_data)
        indices = np.argpartition(dmat, nhood_size)[:, :nhood_size]
        dmat_shortened = submatrix(dmat, indices, nhood_size)
        indices_sorted = np.argsort(dmat_shortened)

# Path: umap/utils.py
# @numba.njit(parallel=True)
# def submatrix(dmat, indices_col, n_neighbors):
#     """Return a submatrix given an orginal matrix and the indices to keep.
# 
#     Parameters
#     ----------
#     dmat: array, shape (n_samples, n_samples)
#         Original matrix.
# 
#     indices_col: array, shape (n_samples, n_neighbors)
#         Indices to keep. Each row consists of the indices of the columns.
# 
#     n_neighbors: int
#         Number of neighbors.
# 
#     Returns
#     -------
#     submat: array, shape (n_samples, n_neighbors)
#         The corresponding submatrix.
#     """
#     n_samples_transform, n_samples_fit = dmat.shape
#     submat = np.zeros((n_samples_transform, n_neighbors), dtype=dmat.dtype)
#     for i in numba.prange(n_samples_transform):
#         for j in numba.prange(n_neighbors):
#             submat[i, j] = dmat[i, indices_col[i, j]]
#     return submat
# 
# def average_nn_distance(dist_matrix):
#     """Calculate the average distance to each points nearest neighbors.
# 
#     Parameters
#     ----------
#     dist_matrix: a csr_matrix
#         A distance matrix (usually umap_model.graph_)
# 
#     Returns
#     -------
#     An array with the average distance to each points nearest neighbors
# 
#     """
#     (row_idx, col_idx, val) = scipy.sparse.find(dist_matrix)
# 
#     # Count/sum is done per row
#     count_non_zero_elems = np.bincount(row_idx)
#     sum_non_zero_elems = np.bincount(row_idx, weights=val)
#     averages = sum_non_zero_elems / count_non_zero_elems
# 
#     if any(np.isnan(averages)):
#         warn(
#             "Embedding contains disconnected vertices which will be ignored."
#             "Use umap.utils.disconnected_vertices() to identify them."
#         )
# 
#     return averages


# Path: umap/plot.py
import numpy as np
import numba
    import pandas as pd
    import datashader as ds
    import datashader.transfer_functions as tf
    import datashader.bundling as bd
    import matplotlib.pyplot as plt
    import colorcet
    import matplotlib.colors
    import matplotlib.cm
    import bokeh.plotting as bpl
    import bokeh.transform as btr
    import holoviews as hv
    import holoviews.operation.datashader as hd
import sklearn.decomposition
import sklearn.cluster
import sklearn.neighbors
from warnings import warn
from matplotlib.patches import Patch
from umap.utils import submatrix, average_nn_distance
from bokeh.plotting import show as show_interactive
from bokeh.plotting import output_file, output_notebook
from bokeh.layouts import column
from bokeh.models import CustomJS, TextInput
from matplotlib.pyplot import show as show_static
from warnings import warn


def _get_metric(umap_object):
    if hasattr(umap_object, "metric"):
        return umap_object.metric
    else:
        # Assume euclidean if no attribute per cuML.UMAP
        return "euclidean"


def _get_metric_kwds(umap_object):
    if hasattr(umap_object, "_metric_kwds"):
        return umap_object._metric_kwds
    else:
        # Assume no keywords exist
        return {}


def _embed_datashader_in_an_axis(datashader_image, ax):
    img_rev = datashader_image.data[::-1]
    mpl_img = np.dstack([_blue(img_rev), _green(img_rev), _red(img_rev)])
    ax.imshow(mpl_img)
    return ax


def _nhood_search(umap_object, nhood_size):
    if hasattr(umap_object, "_small_data") and umap_object._small_data:
        dmat = sklearn.metrics.pairwise_distances(umap_object._raw_data)
        indices = np.argpartition(dmat, nhood_size)[:, :nhood_size]
        dmat_shortened = submatrix(dmat, indices, nhood_size)
        indices_sorted = np.argsort(dmat_shortened)

        indices = submatrix(indices, indices_sorted, nhood_size)