# Copyright 1999-2021 Alibaba Group Holding Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import warnings
from abc import ABCMeta, abstractmethod

import numpy as np
from sklearn.base import BaseEstimator, MultiOutputMixin

from ... import tensor as mt
from ...tensor.reshape.reshape import _reshape as reshape_unchecked
from ..metrics.pairwise import PAIRWISE_DISTANCE_FUNCTIONS
from ..metrics import pairwise_distances_topk
from ..utils import check_array
from ..utils.validation import check_is_fitted
from ._ball_tree import create_ball_tree, ball_tree_query, SklearnBallTree
from ._kd_tree import create_kd_tree, kd_tree_query, SklearnKDTree
from ._faiss import build_faiss_index, faiss_query, METRIC_TO_FAISS_METRIC_TYPE
from ._proxima import build_proxima_index, proxima_query, METRIC_TO_PROXIMA_METRIC_TYPE
from ._kneighbors_graph import KNeighborsGraph


def _call_if_callable(metrics):
    if callable(metrics):
        return metrics()
    return metrics


VALID_METRICS = dict(
    ball_tree=_call_if_callable(SklearnBallTree.valid_metrics),
    kd_tree=_call_if_callable(SklearnKDTree.valid_metrics),
    # The following list comes from the
    # sklearn.metrics.pairwise doc string
    brute=(
        list(PAIRWISE_DISTANCE_FUNCTIONS.keys())
        + [
            "braycurtis",
            "canberra",
            "chebyshev",
            "correlation",
            "cosine",
            "dice",
            "hamming",
            "jaccard",
            "kulsinski",
            "mahalanobis",
            "matching",
            "minkowski",
            "rogerstanimoto",
            "russellrao",
            "seuclidean",
            "sokalmichener",
            "sokalsneath",
            "sqeuclidean",
            "yule",
            "wminkowski",
        ]
    ),
    faiss=list(METRIC_TO_FAISS_METRIC_TYPE),
    proxima=list(METRIC_TO_PROXIMA_METRIC_TYPE),
)


VALID_METRICS_SPARSE = dict(
    ball_tree=[], kd_tree=[], brute=(PAIRWISE_DISTANCE_FUNCTIONS.keys() - {"haversine"})
)


class NeighborsBase(BaseEstimator, MultiOutputMixin, metaclass=ABCMeta):
    """Base class for nearest neighbors estimators."""

    @abstractmethod
    def __init__(
        self,
        n_neighbors=None,
        radius=None,
        algorithm="auto",
        leaf_size=30,
        metric="minkowski",
        p=2,
        metric_params=None,
        n_jobs=None,
    ):
        self.n_neighbors = n_neighbors
        self.radius = radius
        self.algorithm = algorithm
        self.leaf_size = leaf_size
        self.metric = metric
        self.metric_params = metric_params
        self.p = p
        self.n_jobs = n_jobs
        self._check_algorithm_metric()

    def _check_algorithm_metric(self):
        if self.algorithm not in [
            "auto",
            "brute",
            "kd_tree",
            "ball_tree",
            "faiss",
            "proxima",
        ]:
            raise ValueError(f"unrecognized algorithm: '{self.algorithm}'")

        if self.algorithm == "auto":
            if self.metric == "precomputed":
                alg_check = "brute"
            elif callable(self.metric) or self.metric in VALID_METRICS["ball_tree"]:
                alg_check = "ball_tree"
            else:
                alg_check = "brute"
        else:
            alg_check = self.algorithm

        if callable(self.metric):
            if self.algorithm == "kd_tree":
                # callable metric is only valid for brute force and ball_tree
                raise ValueError(
                    "kd_tree algorithm does not support callable metric '%s'"
                    % self.metric
                )
        elif self.metric not in VALID_METRICS[alg_check]:
            raise ValueError(
                "Metric '%s' not valid. Use "
                "sorted(sklearn.neighbors.VALID_METRICS['%s']) "
                "to get valid options. "
                "Metric can also be a callable function." % (self.metric, alg_check)
            )

        if self.metric_params is not None and "p" in self.metric_params:
            warnings.warn(
                "Parameter p is found in metric_params. "
                "The corresponding parameter from __init__ "
                "is ignored.",
                SyntaxWarning,
                stacklevel=3,
            )
            effective_p = self.metric_params["p"]
        else:
            effective_p = self.p

        if self.metric in ["wminkowski", "minkowski"] and effective_p < 1:
            raise ValueError("p must be greater than one for minkowski metric")

    def _fit(self, X, session=None, run_kwargs=None):
        self._check_algorithm_metric()
        if self.metric_params is None:
            self.effective_metric_params_ = {}
        else:
            self.effective_metric_params_ = self.metric_params.copy()

        effective_p = self.effective_metric_params_.get("p", self.p)
        if self.metric in ["wminkowski", "minkowski"]:
            self.effective_metric_params_["p"] = effective_p

        self.effective_metric_ = self.metric
        # For minkowski distance, use more efficient methods where available
        if self.metric == "minkowski":
            p = self.effective_metric_params_.pop("p", 2)
            if p < 1:  # pragma: no cover
                raise ValueError("p must be greater than one for minkowski metric")
            elif p == 1:
                self.effective_metric_ = "manhattan"
            elif p == 2:
                self.effective_metric_ = "euclidean"
            elif p == np.inf:
                self.effective_metric_ = "chebyshev"
            else:
                self.effective_metric_params_["p"] = p

        if isinstance(X, NeighborsBase):
            self._fit_X = X._fit_X
            self._tree = X._tree
            self._fit_method = X._fit_method
            return self

        elif isinstance(X, SklearnBallTree):
            self._fit_X = mt.tensor(X.data)
            self._tree = X
            self._fit_method = "ball_tree"
            return self

        elif isinstance(X, SklearnKDTree):
            self._fit_X = mt.tensor(X.data)
            self._tree = X
            self._fit_method = "kd_tree"
            return self

        X = check_array(X, accept_sparse=True)

        if np.isnan(X.size):
            # if X has unknown shape, execute it first
            X.execute(session=session, **(run_kwargs or dict()))

        if X.issparse():
            if self.algorithm not in ("auto", "brute"):
                warnings.warn("cannot use tree with sparse input: using brute force")
            if self.effective_metric_ not in VALID_METRICS_SPARSE[
                "brute"
            ] and not callable(self.effective_metric_):
                raise ValueError(
                    "Metric '%s' not valid for sparse input. "
                    "Use sorted(sklearn.neighbors."
                    "VALID_METRICS_SPARSE['brute']) "
                    "to get valid options. "
                    "Metric can also be a callable function." % (self.effective_metric_)
                )
            self._fit_X = X.copy()
            self._tree = None
            self._fit_method = "brute"
            return self

        self._fit_method = self.algorithm
        self._fit_X = X

        if self._fit_method == "auto":
            # A tree approach is better for small number of neighbors,
            # and KDTree is generally faster when available
            if (
                self.n_neighbors is None or self.n_neighbors < self._fit_X.shape[0] // 2
            ) and self.metric != "precomputed":
                if self.effective_metric_ in VALID_METRICS["kd_tree"]:
                    self._fit_method = "kd_tree"
                elif (
                    callable(self.effective_metric_)
                    or self.effective_metric_ in VALID_METRICS["ball_tree"]
                ):
                    self._fit_method = "ball_tree"
                else:
                    self._fit_method = "brute"
            else:
                self._fit_method = "brute"

        if self._fit_method == "ball_tree":
            self._tree = tree = create_ball_tree(
                X,
                self.leaf_size,
                metric=self.effective_metric_,
                **self.effective_metric_params_,
            )
            tree.execute(session=session, **(run_kwargs or dict()))
        elif self._fit_method == "kd_tree":
            self._tree = tree = create_kd_tree(
                X,
                self.leaf_size,
                metric=self.effective_metric_,
                **self.effective_metric_params_,
            )
            tree.execute(session=session, **(run_kwargs or dict()))
        elif self._fit_method == "brute":
            self._tree = None
        elif self._fit_method == "faiss":
            faiss_index = build_faiss_index(
                X, metric=self.effective_metric_, **self.effective_metric_params_
            )
            faiss_index.execute(session=session, **(run_kwargs or dict()))
            self._faiss_index = faiss_index
        elif self._fit_method == "proxima":  # pragma: no cover
            proxima_metric = METRIC_TO_PROXIMA_METRIC_TYPE[self.effective_metric_]
            proxima_index = build_proxima_index(
                X,
                distance_metric=proxima_metric,
                topk=self.n_neighbors,
                session=session,
                run_kwargs=run_kwargs,
                **self.effective_metric_params_,
            )
            self._proxima_index = proxima_index
        else:  # pragma: no cover
            raise ValueError("algorithm = '%s' not recognized" % self.algorithm)

        if self.n_neighbors is not None:
            if self.n_neighbors <= 0:
                raise ValueError(f"Expected n_neighbors > 0. Got {self.n_neighbors}")
            else:
                if not np.issubdtype(type(self.n_neighbors), np.integer):
                    raise TypeError(
                        f"n_neighbors does not take {type(self.n_neighbors)} value, "
                        "enter integer value"
                    )

        return self


class KNeighborsMixin:
    """Mixin for k-neighbors searches"""

    def kneighbors(
        self,
        X=None,
        n_neighbors=None,
        return_distance=True,
        session=None,
        run_kwargs=None,
        **kw,
    ):
        """Finds the K-neighbors of a point.
        Returns indices of and distances to the neighbors of each point.

        Parameters
        ----------
        X : array-like, shape (n_query, n_features), \
                or (n_query, n_indexed) if metric == 'precomputed'
            The query point or points.
            If not provided, neighbors of each indexed point are returned.
            In this case, the query point is not considered its own neighbor.

        n_neighbors : int
            Number of neighbors to get (default is the value
            passed to the constructor).

        return_distance : boolean, optional. Defaults to True.
            If False, distances will not be returned

        Returns
        -------
        dist : Tensor
            Array representing the lengths to points, only present if
            return_distance=True

        ind : Tensor
            Indices of the nearest points in the population matrix.

        Examples
        --------
        In the following example, we construct a NeighborsClassifier
        class from a tensor representing our data set and ask who's
        the closest point to [1,1,1]

        >>> samples = [[0., 0., 0.], [0., .5, 0.], [1., 1., .5]]
        >>> from mars.learn.neighbors import NearestNeighbors
        >>> neigh = NearestNeighbors(n_neighbors=1)
        >>> neigh.fit(samples) # doctest: +ELLIPSIS
        NearestNeighbors(algorithm='auto', leaf_size=30, ...)
        >>> print(neigh.kneighbors([[1., 1., 1.]])) # doctest: +ELLIPSIS
        (array([[0.5]]), array([[2]]))

        As you can see, it returns [[0.5]], and [[2]], which means that the
        element is at distance 0.5 and is the third element of samples
        (indexes start at 0). You can also query for multiple points:

        >>> X = [[0., 1., 0.], [1., 0., 1.]]
        >>> neigh.kneighbors(X, return_distance=False) # doctest: +ELLIPSIS
        array([[1],
               [2]]...)

        """
        check_is_fitted(self, ["_fit_method", "_fit_X"], all_or_any=any)

        if n_neighbors is None:
            n_neighbors = self.n_neighbors
        elif n_neighbors <= 0:
            raise ValueError(f"Expected n_neighbors > 0. Got {n_neighbors}")
        else:
            if not np.issubdtype(type(n_neighbors), np.integer):
                raise TypeError(
                    f"n_neighbors does not take {type(n_neighbors)} value, "
                    "enter integer value"
                )

        if X is not None:
            query_is_train = False
            X = check_array(X, accept_sparse=True)
        else:
            query_is_train = True
            X = self._fit_X
            # Include an extra neighbor to account for the sample itself being
            # returned, which is removed later
            n_neighbors += 1

        if X.key == self._fit_X.key and X is not self._fit_X:
            X = self._fit_X
        if np.isnan(X.size):
            # has unknown size, execute first
            X.execute(session=session, **(run_kwargs or dict()))

        train_size = self._fit_X.shape[0]
        if n_neighbors > train_size:
            raise ValueError(
                "Expected n_neighbors <= n_samples, "
                f"but n_samples = {train_size}, n_neighbors = {n_neighbors}"
            )
        n_samples, _ = X.shape
        sample_range = mt.arange(n_samples)[:, None]

        if self._fit_method == "brute":
            # for efficiency, use squared euclidean distances
            kwds = (
                {"squared": True}
                if self.effective_metric_ == "euclidean"
                else self.effective_metric_params_
            )

            neigh_dist, neigh_ind = pairwise_distances_topk(
                X, self._fit_X, k=n_neighbors, metric=self.effective_metric_, **kwds
            )
            if return_distance:
                if self.effective_metric_ == "euclidean":
                    result = mt.sqrt(neigh_dist), neigh_ind
                else:
                    result = neigh_dist, neigh_ind
            else:
                result = neigh_ind
        elif self._fit_method in ["ball_tree", "kd_tree"]:
            if X.issparse():
                raise ValueError(
                    f"{self._fit_method} does not work with sparse matrices. "
                    "Densify the data, or set algorithm='brute'"
                )

            query = (
                ball_tree_query if self._fit_method == "ball_tree" else kd_tree_query
            )
            result = query(self._tree, X, n_neighbors, return_distance)
        elif self._fit_method == "faiss":
            if X.issparse():
                raise ValueError(
                    f"{self._fit_method} does not work with sparse matrices. "
                    "Densify the data, or set algorithm='brute'"
                )
            result = faiss_query(
                self._faiss_index, X, n_neighbors, return_distance, **kw
            )
        elif self._fit_method == "proxima":  # pragma: no cover
            if X.issparse():
                raise ValueError(
                    f"{self._fit_method} does not work with sparse matrices. "
                    "Densify the data, or set algorithm='brute'"
                )
            ind, dis = proxima_query(
                X, n_neighbors, index=self._proxima_index, run=False, **kw
            )
            if not return_distance:
                result = ind
            else:
                result = (dis, ind)
        else:  # pragma: no cover
            raise ValueError("internal: _fit_method not recognized")

        if not query_is_train:
            if isinstance(result, (tuple, list)):
                result = mt.ExecutableTuple(result)
            result.execute(session=session, **(run_kwargs or dict()))
            return result
        else:
            # If the query data is the same as the indexed data, we would like
            # to ignore the first nearest neighbor of every sample, i.e
            # the sample itself.
            if return_distance:
                dist, neigh_ind = result
            else:
                neigh_ind = result

            sample_mask = neigh_ind != sample_range

            # Corner case: When the number of duplicates are more
            # than the number of neighbors, the first NN will not
            # be the sample, but a duplicate.
            # In that case mask the first duplicate.
            dup_gr_nbrs = mt.all(sample_mask, axis=1)
            sample_mask[:, 0] = mt.where(dup_gr_nbrs, False, sample_mask[:, 0])

            neigh_ind = reshape_unchecked(
                neigh_ind[sample_mask], (n_samples, n_neighbors - 1)
            )

            if return_distance:
                dist = reshape_unchecked(
                    dist[sample_mask], (n_samples, n_neighbors - 1)
                )
                ret = mt.ExecutableTuple([dist, neigh_ind])
                ret.execute(session=session, **(run_kwargs or dict()))
                return ret
            neigh_ind.execute(session=session, **(run_kwargs or dict()))
            return neigh_ind

    def kneighbors_graph(
        self,
        X=None,
        n_neighbors=None,
        mode="connectivity",
        session=None,
        run_kwargs=None,
    ):
        """Computes the (weighted) graph of k-Neighbors for points in X

        Parameters
        ----------
        X : array-like, shape (n_query, n_features), \
                or (n_query, n_indexed) if metric == 'precomputed'
            The query point or points.
            If not provided, neighbors of each indexed point are returned.
            In this case, the query point is not considered its own neighbor.

        n_neighbors : int
            Number of neighbors for each sample.
            (default is value passed to the constructor).

        mode : {'connectivity', 'distance'}, optional
            Type of returned matrix: 'connectivity' will return the
            connectivity matrix with ones and zeros, in 'distance' the
            edges are Euclidean distance between points.

        Returns
        -------
        A : SparseTensor, shape = [n_samples, n_samples_fit]
            n_samples_fit is the number of samples in the fitted data
            A[i, j] is assigned the weight of edge that connects i to j.

        Examples
        --------
        >>> X = [[0], [3], [1]]
        >>> from mars.learn.neighbors import NearestNeighbors
        >>> neigh = NearestNeighbors(n_neighbors=2)
        >>> neigh.fit(X) # doctest: +ELLIPSIS
        NearestNeighbors(algorithm='auto', leaf_size=30, ...)
        >>> A = neigh.kneighbors_graph(X)
        >>> A.fetch().toarray()
        array([[1., 0., 1.],
               [0., 1., 1.],
               [1., 0., 1.]])

        See also
        --------
        NearestNeighbors.radius_neighbors_graph
        """
        check_is_fitted(self, ["_fit_method", "_fit_X"], all_or_any=any)
        if n_neighbors is None:
            n_neighbors = self.n_neighbors

        # kneighbors does the None handling.
        if X is not None:
            X = check_array(X, accept_sparse=True)
            n_samples1 = X.shape[0]
        else:
            n_samples1 = self._fit_X.shape[0]

        n_samples2 = self._fit_X.shape[0]

        if mode == "connectivity":
            A_data = None
            A_ind = self.kneighbors(X, n_neighbors, return_distance=False)

        elif mode == "distance":
            A_data, A_ind = self.kneighbors(X, n_neighbors, return_distance=True)

        else:
            raise ValueError(
                'Unsupported mode, must be one of "connectivity" '
                f'or "distance" but got {mode} instead'
            )

        op = KNeighborsGraph(
            a_data=A_data, a_ind=A_ind, n_neighbors=n_neighbors, sparse=True
        )
        graph = op(A_data, A_ind, shape=(n_samples1, n_samples2))
        graph.execute(session=session, **(run_kwargs or dict()))
        return graph


class UnsupervisedMixin:
    def fit(self, X, y=None, session=None, run_kwargs=None):
        """Fit the model using X as training data

        Parameters
        ----------
        X : {array-like, tensor, BallTree, KDTree}
            Training data. If tensor, shape [n_samples, n_features],
            or [n_samples, n_samples] if metric='precomputed'.
        """
        return self._fit(X, session=session, run_kwargs=run_kwargs)
