import logging

import numpy as np
from sklearn.metrics.pairwise import pairwise_distances, 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

logger = logging.getLogger(__name__)


class KMedoids:
    """k-medoids clustering.

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

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

    metric : str, or callable, optional
        What distance metric to use. See :func:metrics.pairwise_distances
        metric can be 'precomputed', the user must then feed the fit method
        with a precomputed kernel matrix and not the design matrix X.
        Default is 'euclidean'. Possible values are 'euclidean', 'manhattan',

    method : str, optional
        Which algorithm to use. 'alternate' is faster while 'pam' is more accurate.
        Default is 'alternate'. Possible values are 'alternate' and 'pam'.

    init : str, optional
        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.
        'build' is a greedy initialization of the medoids used in the original PAM
        algorithm. Often 'build' is more efficient but slower than other
        initializations on big datasets and it is also very non-robust,
        if there are outliers in the dataset, use another initialization.
        Default is 'heuristic'. Possible values are 'random', 'heuristic', 'k-medoids++' and 'build'.

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

    max_iter : int, optional
        Specify the maximum number of iterations when fitting. It can be zero in
        which case only the initialization is computed which may be suitable for
        large datasets when the initialization is sufficiently efficient
        (i.e. for 'build' init). Default is 300.

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

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

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

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

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

    n_iter_ : int
        Number of iterations run.

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

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

    Methods
    -------
    fit(X)
        Compute k-medoids clustering.

    predict(X)
        Predict the closest cluster each sample in X belongs to.

    transform(X)
        Transform X to a cluster-distance space.

    assign()
        Assigns every point in the dataset to the closest center.
    """

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

    def _check_nonnegative_int(self, value, desc, strict=True):
        """
        Check if a value is a nonnegative integer.

        Parameters
        ----------
        value : int
            The value to check.
        desc : str
            The description of the value.
        strict : bool, optional
            If True, the value must be strictly positive. Default is True.

        Raises
        ------
        ValueError
            If the value is not a nonnegative integer.
        """
        negative = value is None or value <= 0 if strict else value is None or value < 0
        if negative or not isinstance(value, (int, np.integer)):
            raise ValueError(f"{desc} should be a nonnegative integer. " f"{value} was given")

    def _check_init_args(self):
        """
        Check the validity of the input arguments.

        Raises
        ------
        ValueError
            If the input arguments are invalid.
        """

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

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

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

        Parameters
        ----------
        X : array-like or sparse matrix
            The data to cluster. It must be an array of shape (n_samples, n_features) or (n_samples, n_samples)
            if metric == 'precomputed'

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

        sample_weight : array-like, shape (n_samples,), optional
            The weights for each sample. If None, all samples are given equal weight.

        Returns
        -------
        self : object
            Returns the instance itself.
        """
        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]:
            raise ValueError(
                "The number of medoids (%d) must be less "
                "than the number of samples %d." % (self.n_clusters, X.shape[0])
            )

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

        # Continue the algorithm as long as
        # the medoids keep changing and the maximum number
        # of iterations is not exceeded

        for n_iter_ in range(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, sample_weight)

            if np.all(old_medoid_idxs == medoid_idxs):
                break
            if n_iter_ == self.max_iter - 1:
                logger.warning(
                    "Maximum number of iteration reached before "
                    "convergence. Consider increasing max_iter to "
                    "improve the fit."
                )

        # Set the resulting instance variables.
        if self.metric == "precomputed":
            self.cluster_centers_ = None
        else:
            self.cluster_centers_ = X[medoid_idxs]
        self.centers = medoid_idxs
        # Expose labels_ which are the assignments of
        # the training data to clusters
        self.labels_ = np.argmin(D[medoid_idxs, :], axis=0)
        self.members = self.labels_
        self.medoid_indices_ = medoid_idxs
        self.inertia_ = self._compute_inertia(self.transform(X))

        # Return self to enable method chaining
        return self

    def assign(self):
        """
        Assigns every point in the dataset to the closest center.

        Returns
        -------
        list
            The indices of the medoid rows in X
        """
        return list(enumerate(self.medoid_indices_))

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

        Parameters
        ----------
        D : array-like
            The distance matrix between the data points, shape (n_samples, n_samples)
        labels : array-like
            The cluster labels for each data point, shape (n_samples,)
        medoid_idxs : array-like
            The indices of the medoid rows in X, shape = (n_clusters,)
        sample_weight : array-like, optional
            The weights for each sample. If None, all samples are given equal weight, shape (n_samples,)
        """

        # 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:
                logger.warning(
                    f"Cluster {k} is empty! "
                    f"self.labels_[self.medoid_indices_[{k}]] "
                    "may not be labeled with "
                    f"its corresponding cluster ({k})."
                )
                continue

            in_cluster_distances = D[cluster_k_idxs, cluster_k_idxs[:, np.newaxis]]

            # Calculate all costs from each point to all others in the cluster
            in_cluster_all_costs = np.sum(in_cluster_distances, axis=1)

            if sample_weight is not None:
                in_cluster_all_costs *= sample_weight[cluster_k_idxs]

            min_cost_idx = np.argmin(in_cluster_all_costs)
            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 then the current
            if min_cost < curr_cost:
                medoid_idxs[k] = cluster_k_idxs[min_cost_idx]

    def _compute_cost(self, D, medoid_idxs):
        """
        Compute the cose for a given configuration of the medoids.

        Parameters
        ----------
        D : array-like
            The distance matrix between the data points, shape (n_samples, n_samples)
        medoid_idxs : array-like
            The indices of the medoid rows in X, shape = (n_clusters,)

        Returns
        -------
        float
            The cost of the configuration.
        """
        return self._compute_inertia(D[:, medoid_idxs])

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

        Parameters
        ----------
        X : array-like or sparse matrix
            Data to transform. It must be an array of shape (n_query, n_features) or (n_query, n_indexed)
            if metric == 'precomputed'

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

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

        check_is_fitted(self, "cluster_centers_")

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

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

        Parameters
        ----------
        X : array-like or sparse matrix
            New data to predict. It must be an array of shape (n_query, n_features) or (n_query, n_indexed)
            if metric == 'precomputed'

        Returns
        -------
        array-like
            Index of the cluster each sample belongs to, shape = (n_query,)
        """
        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)

        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 or sparse matrix
            Distances to cluster centers, shape=(n_samples, n_clusters)

        Returns
        -------
        float
            Sum of distances of samples to their closest cluster center.
        """

        # 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_):
        """
        Select initial mediods when beginning clustering.

        Parameters
        ----------
        D : array-like
            The distance matrix between the data points, shape (n_samples, n_samples)
        n_clusters : int
            The number of clusters to form as well as the number of medoids to generate.
        random_state_ : RandomState instance
            The random number generator.

        Returns
        -------
        array-like
            The indices of the medoid rows in X, shape = (n_clusters,)
        """

        if self.init == "random":  # Random initialization
            # Pick random k medoids as the initial ones.
            medoids = random_state_.choice(len(D), n_clusters)
        elif self.init == "k-medoids++":
            medoids = self._kpp_init(D, n_clusters, random_state_)
        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.
            medoids = np.argpartition(np.sum(D, axis=1), n_clusters - 1)[:n_clusters]
        else:
            raise ValueError(f"init value '{self.init}' not recognized")

        return medoids

    def _kpp_init(self, D, n_clusters, random_state_, n_local_trials=None):
        """
        Init n_clusters seeds with a method similar to k-means++

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

        n_clusters : int
            The number of seeds to choose

        random_state : RandomState instance
            The generator used to initialize the centers.

        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.

        Returns
        -------
        array-like
            The indices of the medoid rows in X, shape = (n_clusters,)

        References
        ----------
        .. [1] Arthur, D. and Vassilvitskii, S.
        "k-means++: the advantages of careful seeding". ACM-SIAM symposium
        on Discrete algorithms. 2007

        Notes
        -----
        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))

        center_id = random_state_.randint(n_samples)
        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):
            rand_vals = random_state_.random_sample(n_local_trials) * current_pot
            candidate_ids = np.searchsorted(stable_cumsum(closest_dist_sq), rand_vals)

            # 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
