# -*- coding: utf-8 -*-
# @Author: Aliaksei Halachkin
# @Date:   2017-01-15 18:02:06
# @Last Modified by:   Aliaksei Halachkin
# @Last Modified time: 2017-04-19 15:06:55

from __future__ import division, print_function

from collections import Counter

import numpy as np
from sklearn.cluster import KMeans
from sklearn.utils import check_random_state
from imblearn.base import BaseMulticlassSampler


class ClusterCentroidsMod(BaseMulticlassSampler):
    """
    Modified imbalancedlearn ClusterCentroids class
    n_clusters could be float or integer
    """

    def __init__(self,
                 n_clusters=None,
                 random_state=None,
                 estimator=None,
                 n_jobs=1):
        super(ClusterCentroidsMod, self).__init__(
            ratio='auto', random_state=random_state)
        self.estimator = estimator
        self.n_jobs = n_jobs
        self.n_clusters = n_clusters

    def _validate_estimator(self):
        """Private function to create the NN estimator"""

        if self.estimator is None:
            self.estimator_ = KMeans(
                random_state=self.random_state, n_jobs=self.n_jobs)
        elif isinstance(self.estimator, KMeans):
            self.estimator_ = self.estimator
        else:
            raise ValueError('`estimator` has to be a KMeans clustering.')

    def fit(self, X, y):
        """Find the classes statistics before to perform sampling.

        Parameters
        ----------
        X : ndarray, shape (n_samples, n_features)
            Matrix containing the data which have to be sampled.

        y : ndarray, shape (n_samples, )
            Corresponding label for each sample in X.

        Returns
        -------
        self : object,
            Return self.

        """

        super(ClusterCentroidsMod, self).fit(X, y)

        self._validate_estimator()

        return self

    def _sample(self, X, y):
        """Resample the dataset.

        Parameters
        ----------
        X : ndarray, shape (n_samples, n_features)
            Matrix containing the data which have to be sampled.

        y : ndarray, shape (n_samples, )
            Corresponding label for each sample in X.

        Returns
        -------
        X_resampled : ndarray, shape (n_samples_new, n_features)
            The array containing the resampled data.

        y_resampled : ndarray, shape (n_samples_new)
            The corresponding label of `X_resampled`

        """

        X_resampled = None
        y_resampled = None

        
        for key in self.stats_c_.keys():
            if type(self.n_clusters) is int:
                n_clusters = self.n_clusters
            elif type(self.n_clusters) is float:
                n_clusters = int(np.ceil(self.n_clusters * self.stats_c_[key]))
            else:
                raise ValueError("n_clusters should be float or int")

            if n_clusters >= self.stats_c_[key]:
                centroids = X[y==key]
                n_clusters = self.stats_c_[key]
            else:
                self.estimator_.set_params(**{'n_clusters': n_clusters})
                self.estimator_.fit(X[y == key])
                centroids = self.estimator_.cluster_centers_

            if X_resampled is None and y_resampled is None:
                X_resampled = centroids
                y_resampled = np.array([key] * n_clusters)
            else:
                X_resampled = np.concatenate((X_resampled, centroids), axis=0)
                y_resampled = np.concatenate(
                    (y_resampled, np.array([key] * n_clusters)), axis=0)

        self.logger.info('Under-sampling performed: %s', Counter(y_resampled))

        return X_resampled, y_resampled


class RandomUnderSamplerMod(BaseMulticlassSampler):
    """Class to perform random under-sampling.

    """

    def __init__(self,
                 num_samples,
                 # return_indices=False,
                 random_state=None,
                 replacement=True):
        super(RandomUnderSamplerMod, self).__init__(
            ratio='auto', random_state=random_state)
        # self.return_indices = return_indices
        self.replacement = replacement
        self.num_samples = num_samples

    def _sample(self, X, y):
        """Resample the dataset.

        Parameters
        ----------
        X : ndarray, shape (n_samples, n_features)
            Matrix containing the data to be sampled.

        y : ndarray, shape (n_samples, )
            Corresponding label for each sample in X.

        Returns
        -------
        X_resampled : ndarray, shape (n_samples_new, n_features)
            The array containing the resampled data.

        y_resampled : ndarray, shape (n_samples_new)
            The corresponding label of `X_resampled`

        idx_under : ndarray, shape (n_samples, )
            If `return_indices` is `True`, an array will be returned
            containing a boolean for each sample to represent whether
            that sample was selected or not.

        """

        random_state = check_random_state(self.random_state)

        # Compute the number of clusters needed
        # if self.ratio == 'auto':
        #     num_samples = self.stats_c_[self.min_c_]
        # else:
        #     num_samples = int(self.stats_c_[self.min_c_] / self.ratio)

        # All the minority class samples will be preserved
        # X_resampled = X[y == self.min_c_]
        # y_resampled = y[y == self.min_c_]

        X_resampled = None
        y_resampled = None

        # If we need to offer support for the indices
        # if self.return_indices:
        #     idx_under = np.nonzero(y == self.min_c_)[0]

        # Loop over the other classes under-picking at random
        for key in self.stats_c_.keys():

            # If the minority class is up, skip it
            # if key == self.min_c_:
            #     continue

            if type(self.num_samples) is int:
                num_samples = self.num_samples
                if num_samples > self.stats_c_[key]:
                    num_samples = self.stats_c_[key]
            elif type(self.num_samples) is float:
                num_samples = int(np.ceil(self.stats_c_[key] * self.num_samples))
            else:
                raise ValueError("num_samples should be float or int")

            # Pick some elements at random
            indx = range(np.count_nonzero(y == key))
            indx = random_state.choice(
                indx, size=num_samples, replace=self.replacement)

            # If we need to offer support for the indices selected
            # if self.return_indices:
            #     idx_tmp = np.nonzero(y == key)[0][indx]
            #     idx_under = np.concatenate((idx_under, idx_tmp), axis=0)

            # Concatenate to the minority class
            if X_resampled is None and y_resampled is None:
                X_resampled = X[y == key][indx]
                y_resampled = y[y == key][indx]
            else:
                X_resampled = np.concatenate(
                    (X_resampled, X[y == key][indx]), axis=0)
                y_resampled = np.concatenate(
                    (y_resampled, y[y == key][indx]), axis=0)

        self.logger.info('Under-sampling performed: %s', Counter(y_resampled))

        # Check if the indices of the samples selected should be returned as
        # well
        # if self.return_indices:
        #     # Return the indices of interest
        #     return X_resampled, y_resampled, idx_under
        # else:
        #     return X_resampled, y_resampled

        return X_resampled, y_resampled
