"""K-medoids clustering"""
import warnings
import numpy as np
from sklearn.base import BaseEstimator, ClusterMixin, TransformerMixin
from sklearn.metrics.pairwise import pairwise_distances
from sklearn.metrics.pairwise import pairwise_distances_argmin
from sklearn.utils import check_array, check_random_state
from sklearn.utils.extmath import stable_cumsum
from sklearn.utils.validation import check_is_fitted, _check_sample_weight
from sklearn.utils.validation import _num_samples
from sklearn.exceptions import ConvergenceWarning
from scipy.spatial import distance
from scipy.spatial.distance import _validate_vector
from scipy.linalg import norm


def abs_correlation(u, v):
    """
    absolute pearson correlation distance (1-abs(pearson_cc))
    """
    u = _validate_vector(u)
    v = _validate_vector(v)
    umu = u.mean()
    vmu = v.mean()
    um = u - umu
    vm = v - vmu
    dist = 1.0 - np.abs(np.dot(um, vm) / (norm(um) * norm(vm)))
    return dist


class KMedoids(BaseEstimator, ClusterMixin, TransformerMixin):
    """weighted k-medoids clustering.

    Read more in the :ref:`User Guide <k_medoids>`.

    Parameters
    ----------
    n_clusters : int, optional, default: 8
        The number of clusters to form as well as the number of medoids to
        generate.

    metric : string, or callable, optional, default: 'euclidean'
        What distance metric to use. See :func:metrics.pairwise_distances

    init : {'random', 'heuristic', 'k-medoids++'}, optional,
        default: 'heuristic'
        Specify medoid initialization method. 'random' selects n_clusters
        elements from the dataset. 'heuristic' picks the n_clusters points
        with the smallest sum distance to every other point. 'k-medoids++'
        follows an approach based on k-means++_, and in general, gives initial
        medoids which are more separated than those generated by the other
        methods.

        .. _k-means++:
        https://theory.stanford.edu/~sergei/papers/kMeansPP-soda.pdf

    max_iter : int, optional, default : 300
        Specify the maximum number of iterations when fitting.

    random_state : int, RandomState instance or None, optional
        Specify random state for the random number generator. Used to
        initialise medoids when init='random'.

    w: array, shape = (n_samples,), optional, default: None
       Specify the sample weights

    Attributes
    ----------
    cluster_centers_ : array, shape = (n_clusters, n_features)
            or None if metric == 'precomputed'
        Cluster centers, i.e. medoids (elements from the original dataset)

    medoid_indices_ : array, shape = (n_clusters,)
        The indices of the medoid rows in X

    labels_ : array, shape = (n_samples,)
        Labels of each point

    inertia_ : float
        Sum of distances of samples to their closest cluster center.

    Examples
    --------
    >>> import numpy as np

    >>> X = np.asarray([[1, 2], [1, 4], [1, 0],
    ...                 [4, 2], [4, 4], [4, 0]])
    >>> kmedoids = KMedoids(n_clusters=2, random_state=0).fit(X)
    >>> kmedoids.labels_
    array([0, 0, 0, 1, 1, 1])
    >>> kmedoids.predict([[0,0], [4,4]])
    array([0, 1])
    >>> kmedoids.cluster_centers_
    array([[1, 2],
           [4, 2]])
    >>> kmedoids.inertia_
    8.0

    See scikit-learn-extra/examples/plot_kmedoids_digits.py for examples
    of KMedoids with various distance metrics.

    References
    ----------
    Kaufman, L. and Rousseeuw, P.J., Statistical Data Analysis Based on
    the L1–Norm and Related Methods, edited by Y. Dodge, North-Holland,
    405–416. 1987

    See also
    --------

    KMeans
        The KMeans algorithm minimizes the within-cluster sum-of-squares
        criterion. It scales well to large number of samples.

    Notes
    -----
    Since all pairwise distances are calculated and stored in memory for
    the duration of fit, the space complexity is O(n_samples ** 2).

    """

    def __init__(
        self,
        n_clusters=8,
        metric="euclidean",
        init="heuristic",
        max_iter=300,
        n_jobs=None,
        random_state=None,
    ):
        self.n_clusters = n_clusters
        self.metric = metric
        self.init = init
        self.max_iter = max_iter
        self.n_jobs = n_jobs
        self.random_state = random_state

    def _check_nonnegative_int(self, value, desc):
        """Validates if value is a valid integer > 0"""

        if (
            value is None
            or value <= 0
            or not isinstance(value, (int, np.integer))
        ):
            raise ValueError(
                "%s should be a nonnegative integer. "
                "%s was given" % (desc, value)
            )

    def _check_normalized_sample_weight(self, sample_weight, X):
        """Set sample_weight to 1 if None
           if provided, normalize them to within the range of (0,1)
        """

        sample_weight_was_none = sample_weight is None
        sample_weight = _check_sample_weight(sample_weight, X, dtype=X.dtype)
        # require postive weight only
        assert np.all(sample_weight > 0), 'weight must be positive'
        # if weights are in [0, 1], do nothing
        if not sample_weight_was_none:
            all_in_range = ((sample_weight >= 0).all() and
                            (sample_weight <= 1).all())
            if not all_in_range:
                # 2*(1-sigmoid(x)) -> (0,1), normalized to between 0 and 1
                sample_weight = 2.0 * (1 / (1 + np.exp(-sample_weight)) - 0.5)
        return sample_weight

    def _check_candidacy(self, candidacy, X):
        """"check candidacy array values"""

        n_samples = _num_samples(X)

        if candidacy is None:
            candidacy = np.ones(n_samples, dtype=bool)
        else:
            if candidacy.ndim != 1:
                raise ValueError('candidacy must be 1D array')

            if candidacy.shape != (n_samples,):
                raise ValueError('candidacy.shape == {}, expected {}!'
                                 .format(candidacy.shape, (n_samples,)))

            if np.issubdtype(candidacy.dtype, np.integer):
                candidacy = candidacy.astype(np.bool)
            elif candidacy.dtype != np.bool:
                raise TypeError('candidacy type must be integer or bool')

        return candidacy

    def _check_init_args(self):
        """Validates the input arguments. """

        # Check n_clusters and max_iter
        self._check_nonnegative_int(self.n_clusters, "n_clusters")
        self._check_nonnegative_int(self.max_iter, "max_iter")

        # Check init
        init_methods = ["random", "heuristic", "k-medoids++"]
        if self.init not in init_methods:
            raise ValueError(
                "init needs to be one of "
                + "the following: "
                + "%s" % init_methods
            )

    def fit(self, X, y=None, sample_weight=None, candidacy=None):
        """Fit K-Medoids to the provided data.

        Parameters
        ----------
        X : {array-like, sparse matrix}, shape = (n_samples, n_features), \
                or (n_samples, n_samples) if metric == 'precomputed'
            Dataset to cluster.

        y : Ignored

        sample_weight : array-like, shape (n_samples,), optional
            The weights for each observation in X. If None, all observations
            are assigned equal weight (default: None).

        candidacy:  array-like, shape (n_sample,), optional
           Only the samples indicated as True in the candidacy can be selected
           as cluster center, if candidacy is None,  any sample can be
           potentially selected as center.

        Returns
        -------
        self
        """
        random_state_ = check_random_state(self.random_state)

        self._check_init_args()
        X = check_array(X, accept_sparse=["csr", "csc"])
        if self.n_clusters > X.shape[0]:
            warn_msg = f'The number of medoids ({self.n_clusters}) must be less ' \
                       f'than the number of samples ({X.shape[0]})'
            warnings.warn(warn_msg)

        cnsw = self._check_normalized_sample_weight
        checked_sample_weight = cnsw(sample_weight, X)
        checked_candidacy = self._check_candidacy(candidacy, X)

        n_cand = np.sum(checked_candidacy)
        if n_cand < self.n_clusters:
            warn_msg = f'The number of candidancy {n_cand} cannot be less '   \
                       f'than the number of clusters {self.n_clusters}'
            warnings.warn(warn_msg)

        D = pairwise_distances(X, metric=self.metric, n_jobs=self.n_jobs)
        medoid_idxs = self._initialize_medoids(D, self.n_clusters,
                                               random_state_,
                                               checked_candidacy)
        labels = None

        # Continue the algorithm as long as
        # the medoids keep changing and the maximum number
        # of iterations is not exceeded
        for self.n_iter_ in range(0, self.max_iter):
            old_medoid_idxs = np.copy(medoid_idxs)
            labels = np.argmin(D[medoid_idxs, :], axis=0)

            # Update medoids with the new cluster indices
            self._update_medoid_idxs_in_place(D, labels, medoid_idxs,
                                              checked_sample_weight,
                                              checked_candidacy)
            if np.all(old_medoid_idxs == medoid_idxs):
                break
            elif self.n_iter_ == self.max_iter - 1:
                warnings.warn(
                    "Maximum number of iteration reached before "
                    "convergence. Consider increasing max_iter to "
                    "improve the fit.",
                    ConvergenceWarning,
                )

        # Set the resulting instance variables.
        if self.metric == "precomputed":
            self.cluster_centers_ = None
        else:
            self.cluster_centers_ = X[medoid_idxs]

        # Expose labels_ which are the assignments of
        # the training data to clusters
        self.labels_ = labels
        self.medoid_indices_ = medoid_idxs
        self.inertia_ = self._compute_inertia(self.transform(X))

        # Return self to enable method chaining
        return self

    def _update_medoid_idxs_in_place(self, D, labels, medoid_idxs,
                                     checked_sample_weight, checked_candidacy):
        """In-place update of the medoid indices"""

        # Update the medoids for each cluster
        for k in range(self.n_clusters):
            # Extract the distance matrix between the data points
            # inside the cluster k
            cluster_k_idxs = np.where(labels == k)[0]

            if len(cluster_k_idxs) == 0:
                warnings.warn(
                    "Cluster {k} is empty! "
                    "self.labels_[self.medoid_indices_[{k}]] "
                    "may not be labeled with "
                    "its corresponding cluster ({k}).".format(k=k)
                )
                continue

            in_cluster_distances = D[np.ix_(cluster_k_idxs, cluster_k_idxs)]
            in_cluster_candidacy = checked_candidacy[cluster_k_idxs]

            # Calculate all costs from each point to all others in
            # the cluster (weighted)
            in_cluster_w = checked_sample_weight[cluster_k_idxs]
            in_cluster_all_costs = in_cluster_distances.dot(in_cluster_w)

            # select the one with min cost among all candidates
            order = np.argsort(in_cluster_all_costs)
            min_cost_idx = order[in_cluster_candidacy[order]][0]
            min_cost = in_cluster_all_costs[min_cost_idx]
            curr_cost = in_cluster_all_costs[np.argmax(cluster_k_idxs ==
                                                       medoid_idxs[k])]

            # Adopt a new medoid if its distance is smaller than the current
            if min_cost < curr_cost:
                medoid_idxs[k] = cluster_k_idxs[min_cost_idx]

    def fit_predict(self, X, y=None, sample_weight=None, candidacy=None):
        """Compute cluster centers and predict cluster index for each sample.

        Convenience method; equivalent to calling fit(X) followed by
        predict(X).

        Parameters
        ----------
        X : {array-like, sparse matrix} of shape (n_samples, n_features)
            New data to transform.

        y : Ignored
            Not used, present here for API consistency by convention.

        sample_weight : array-like, shape (n_samples,), optional
            The weights for each observation in X. If None, all observations
            are assigned equal weight (default: None).

        candidacy:  array-like, shape (n_sample,), optional
           Only the samples indicated as True in the candidacy can be selected
           as cluster center, if candidacy is None,  any sample can be
           potentially
           selected as center.

        Returns
        -------
        labels : array, shape [n_samples,]
            Index of the cluster each sample belongs to.
        """
        return self.fit(X, sample_weight=sample_weight,
                        candidacy=candidacy).labels_

    def transform(self, X):
        """Transforms X to cluster-distance space.

        Parameters
        ----------
        X : {array-like, sparse matrix}, shape (n_query, n_features), \
                or (n_query, n_indexed) if metric == 'precomputed'
            Data to transform.

        Returns
        -------
        X_new : {array-like, sparse matrix}, shape=(n_query, n_clusters)
            X transformed in the new space of distances to cluster centers.
        """
        X = check_array(X, accept_sparse=["csr", "csc"])

        if self.metric == "precomputed":
            check_is_fitted(self, "medoid_indices_")
            return X[:, self.medoid_indices_]
        else:
            check_is_fitted(self, "cluster_centers_")

            Y = self.cluster_centers_
            return pairwise_distances(X, Y=Y, metric=self.metric,
                                      n_jobs=self.n_jobs)

    def predict(self, X):
        """Predict the closest cluster for each sample in X.

        Parameters
        ----------
        X : {array-like, sparse matrix}, shape (n_query, n_features), \
                or (n_query, n_indexed) if metric == 'precomputed'
            New data to predict.

        Returns
        -------
        labels : array, shape = (n_query,)
            Index of the cluster each sample belongs to.
        """
        X = check_array(X, accept_sparse=["csr", "csc"])

        if self.metric == "precomputed":
            check_is_fitted(self, "medoid_indices_")
            return np.argmin(X[:, self.medoid_indices_], axis=1)
        else:
            check_is_fitted(self, "cluster_centers_")

            # Return data points to clusters based on which cluster assignment
            # yields the smallest distance
            return pairwise_distances_argmin(
                X, Y=self.cluster_centers_, metric=self.metric
            )

    def _compute_inertia(self, distances):
        """Compute inertia of new samples. Inertia is defined as the sum of the
        sample distances to closest cluster centers.

        Parameters
        ----------
        distances : {array-like, sparse matrix}, shape=(n_samples, n_clusters)
            Distances to cluster centers.

        Returns
        -------
        Sum of sample distances to closest cluster centers.
        """

        # Define inertia as the sum of the sample-distances
        # to closest cluster centers
        inertia = np.sum(np.min(distances, axis=1))

        return inertia

    def _initialize_medoids(self, D, n_clusters, random_state_,
                            checked_candidacy):
        """Select initial mediods when beginning clustering."""

        if self.init == "random":  # Random initialization
            # Pick random k medoids as the initial ones from candidacies
            candidacy = np.where(checked_candidacy)[0]
            medoids = random_state_.choice(candidacy, n_clusters)
        elif self.init == "k-medoids++":
            medoids = self._kpp_init(D, n_clusters, random_state_,
                                     checked_candidacy)
        elif self.init == "heuristic":  # Initialization by heuristic
            # Pick K first data points that have the smallest sum distance
            # to every other point. These are the initial medoids.
            # sort first
            order = np.argsort(np.sum(D, axis=1))
            medoids = order[checked_candidacy[order]][:n_clusters]
        else:
            raise ValueError(
                "init value '{init}' not recognized".format(init=self.init)
            )

        return medoids

    # Copied from sklearn.cluster.k_means_._k_init
    def _kpp_init(self, D, n_clusters, random_state_, checked_candidacy,
                  n_local_trials=None):
        """Init n_clusters seeds with a method similar to k-means++

        Parameters
        -----------
        D : array, shape (n_samples, n_samples)
            The distance matrix we will use to select medoid indices.

        n_clusters : integer
            The number of seeds to choose

        random_state : RandomState
            The generator used to initialize the centers.

        checked_candidacy:

        n_local_trials : integer, optional
            The number of seeding trials for each center (except the first),
            of which the one reducing inertia the most is greedily chosen.
            Set to None to make the number of trials depend logarithmically
            on the number of seeds (2+log(k)); this is the default.

        Notes
        -----
        Selects initial cluster centers for k-medoid clustering in a smart way
        to speed up convergence. see: Arthur, D. and Vassilvitskii, S.
        "k-means++: the advantages of careful seeding". ACM-SIAM symposium
        on Discrete algorithms. 2007

        Version ported from http://www.stanford.edu/~darthur/kMeansppTest.zip,
        which is the implementation used in the aforementioned paper.
        """
        n_samples, _ = D.shape

        centers = np.empty(n_clusters, dtype=int)

        # Set the number of local seeding trials if none is given
        if n_local_trials is None:
            # This is what Arthur/Vassilvitskii tried, but did not report
            # specific results for other than mentioning in the conclusion
            # that it helped.
            n_local_trials = 2 + int(np.log(n_clusters))

        # select first center from candidates
        center_id = random_state_.choice(np.where(checked_candidacy)[0], 1)
        centers[0] = center_id

        # Initialize list of closest distances and calculate current potential
        closest_dist_sq = D[centers[0], :] ** 2
        current_pot = closest_dist_sq.sum()

        # pick the remaining n_clusters-1 points
        for cluster_index in range(1, n_clusters):
            # Choose center candidates by sampling with prob. proportional
            # to the squared distance to the closest existing center
            candidate_ids = []
            # possibly never ending ??
            while len(candidate_ids) != n_local_trials:
                rand_val = (random_state_.random_sample(1) * current_pot)
                candidate_id = np.searchsorted(stable_cumsum(closest_dist_sq),
                                               rand_val)
                # XXX: numerical imprecision can result in a candidate_id out
                # of range
                candidate_id = np.clip(candidate_id, None,
                                       closest_dist_sq.size - 1)
                if checked_candidacy[candidate_id]:
                    candidate_ids.append(candidate_id)

            assert len(candidate_ids) == n_local_trials

            # Compute distances to center candidates
            distance_to_candidates = D[candidate_ids, :] ** 2

            # Decide which candidate is the best
            best_candidate = None
            best_pot = None
            best_dist_sq = None
            for trial in range(n_local_trials):
                # Compute potential when including center candidate
                new_dist_sq = np.minimum(
                    closest_dist_sq, distance_to_candidates[trial]
                )
                new_pot = new_dist_sq.sum()

                # Store result if it is the best local trial so far
                if (best_candidate is None) or (new_pot < best_pot):
                    best_candidate = candidate_ids[trial]
                    best_pot = new_pot
                    best_dist_sq = new_dist_sq

            centers[cluster_index] = best_candidate
            current_pot = best_pot
            closest_dist_sq = best_dist_sq

        return centers
