import numpy as np
from scipy import sparse as sp
from ..utils import check_random_state
from ..utils.sparsefuncs import mean_variance_axis
from ..utils.validation import check_is_fitted
from ..cluster.k_means_ import k_means
from ..externals.six import string_types
from .utils import getArrayFromNumericTable, make_csr_table, make_HomogenNumericTable
from daal.data_management import CSRNumericTable
from daal.data_management import HomogenNumericTable,HomogenNumericTable_Float32,HomogenNumericTable_Intc
from daal.algorithms.kmeans import Batch_Float64LloydDense, Batch_Float64LloydCSR, init, data, inputCentroids, assignments, centroids, goalFunction, nIterations, Batch_Float32LloydDense, Batch_Float32LloydCSR

def daal_k_means_dense(X, nClusters, numIterations, tol, cluster_centers_0, random_state):

    def is_string(m, target_):
        return isinstance(m, string_types) and m == target_

    def make_algo_token(dt, algo_name, csr):
        """Constructs init algorithm name in DAAL"""
        if dt == np.float64:
            _type = 'Float64'
        elif dt == np.float32:
            _type = 'Float32'
        else:
            raise NotImplementedError('Unsupported input type %s' % dt)
        csr_suffix = 'CSR' if csr else 'Dense'
        return 'Batch_' + _type + algo_name + csr_suffix

    def _compute_centroids_named(initAlg_callable, nClusters, dataTable):
        """Given callable to initialization algorithm, compute initial cluster centers"""
        if initAlg_callable is None:
            raise ValueError('Failed to determine initialization method')
        initAlg = initAlg_callable(nClusters)
        initAlg.input.set(init.data, dataTable)
        res = initAlg.compute()
        centroids_ = res.get(init.centroids)
        return centroids_

    def _compute_starting_centroids(dataTable, cluster_centers_0, nClusters, dt, csr=False):
        """Process initial clusters argument of KMeans, and either compute them using given algorithm"""
        """ or consume explicitly given centers"""
        if is_string(cluster_centers_0, 'k-means++'):
            method_token = make_algo_token(dt, 'PlusPlus', csr)
            initAlg_callable = getattr(init, method_token, None)
            centroids_ = _compute_centroids_named(initAlg_callable, nClusters, dataTable)
        elif is_string(cluster_centers_0, 'random'):
            method_token = make_algo_token(dt, 'Random', csr)
            initAlg_callable = getattr(init, method_token, None)
            centroids_ = _compute_centroids_named(initAlg_callable, nClusters, dataTable)
        elif hasattr(cluster_centers_0, '__array__'):
            cc_arr = np.ascontiguousarray(cluster_centers_0, dtype=X.dtype)
            _validate_center_shape(X, nClusters, cc_arr)
            centroids_ = make_HomogenNumericTable(cc_arr)
        elif callable(cluster_centers_0):
            cc_arr = cluster_centers_0(X, nClusters, random_state)
            cc_arr = np.ascontiguousarray(cc_arr, dtype=X.dtype)
            _validate_center_shape(X, nClusters, cc_arr)
            centroids_ = make_HomogenNumericTable(cc_arr)
        elif is_string(cluster_centers_0, 'deterministic'):
            method_token = make_algo_token(dt, 'Random', csr)
            initAlg_callable = getattr(init, method_token, None)
            centroids_ = _compute_centroids_named(initAlg_callable, nClusters, dataTable)
        else:
            raise ValueError("Cluster centers should either be 'k-means++', 'random', 'deterministic' or an array")
        return centroids_

    def _tolerance(X, rtol):
        """Compute absolute tolerance from the relative tolerance"""
        if rtol == 0.0:
            return rtol
        if sp.issparse(X):
            variances = mean_variance_axis(X, axis=0)[1]
        else:
            variances = np.var(X, axis=0)
        return np.mean(variances) * rtol

    if numIterations < 0:
        raise ValueError("Wrong iterations number")

    if hasattr(X, '__array__'):
        dataTable = make_HomogenNumericTable(X)
        centroids_result = _compute_starting_centroids(dataTable, cluster_centers_0, nClusters, X.dtype, csr=False)

        algorithm = Batch_Float32LloydDense(nClusters, numIterations) if X.dtype == np.float32 else Batch_Float64LloydDense(nClusters, numIterations)

    elif sp.isspmatrix_csr(X):
        dataTable = make_csr_table(X)
        centroids_result = _compute_starting_centroids(dataTable, cluster_centers_0, nClusters, X.dtype, csr=True)

        algorithm = Batch_Float32LloydCSR(nClusters, numIterations) if X.dtype == np.float32 else Batch_Float64LloydCSR(nClusters, numIterations)

    else:
        raise NotImplementedError("unsupported input type. you can disable pydaal with sklearn.daal4sklearn.dispatcher.disable()")

    algorithm.parameter.maxIterations = numIterations
    algorithm.parameter.accuracyThreshold = _tolerance(X, tol)
    algorithm.parameter.assignFlag = True
    algorithm.input.set(data, dataTable)
    algorithm.input.set(inputCentroids, centroids_result)
    res = algorithm.compute()

    centroids_dtype = np.float32 if X.dtype == np.float32 else np.float64
    labels = getArrayFromNumericTable(res.get(assignments))[:,0].astype(np.intc)
    n_iter = int(getArrayFromNumericTable(res.get(nIterations))[0,0])
    # Per documentation, with numIterations == 0, centroids and goalFunction are not updated
    if numIterations > 0:
        cluster_centers = getArrayFromNumericTable(res.get(centroids)).astype(centroids_dtype, copy = False)
        inertia = getArrayFromNumericTable(res.get(goalFunction))[0,0]
    else:
        cluster_centers = None
        inertia = None
    return cluster_centers, labels, inertia, n_iter


def _validate_center_shape(X, n_centers, centers):
    """Check if centers is compatible with X and n_centers"""
    if len(centers) != n_centers:
        raise ValueError('The shape of the initial centers (%s) '
                         'does not match the number of clusters %i'
                         % (centers.shape, n_centers))
    if centers.shape[1] != X.shape[1]:
        raise ValueError(
            "The number of features of the initial centers %s "
            "does not match the number of features of the data %s."
            % (centers.shape[1], X.shape[1]))


def fit(self, x, y=None):
    """compute k-means clustering.

    parameters
    ----------
    x : array-like or sparse matrix, shape=(n_samples, n_features)
    """
    if self.n_init <= 0:
        raise ValueError("Invalid number of initializations."
                         " n_init=%d must be bigger than zero." % self.n_init)

    random_state = check_random_state(self.random_state)
    x = self._check_fit_data(x)

    if self.max_iter <= 0:
        raise ValueError('Number of iterations should be a positive number,'
                         ' got %d instead' % self.max_iter)

    if self.precompute_distances == 'auto':
        precompute_distances = False
    elif isinstance(self.precompute_distances, bool):
        precompute_distances = self.precompute_distances
    else:
        raise ValueError("precompute_distances should be 'auto' or True/False"
                         ", but a value of %r was passed" %
                         self.precompute_distances)

    daal_ready = hasattr(x, '__array__') or sp.isspmatrix_csr(x)
    daal_ready = daal_ready and not precompute_distances

    if not daal_ready:
        self.cluster_centers_, self.labels_, self.inertia_, self.n_iter_ = \
            k_means(
                x, n_clusters=self.n_clusters, init=self.init,
                n_init=self.n_init, max_iter=self.max_iter, verbose=self.verbose,
                precompute_distances=precompute_distances,
                tol=self.tol, random_state=random_state, copy_x=self.copy_x,
                n_jobs=self.n_jobs, algorithm=self.algorithm,
                return_n_iter=True)
    else:
        self.cluster_centers_, self.labels_, self.inertia_, self.n_iter_ = \
            daal_k_means_dense(
                x, self.n_clusters, self.max_iter, self.tol, self.init,
                random_state)
    return self


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

    In the vector quantization literature, `cluster_centers_` is called
    the code book and each value returned by `predict` is the index of
    the closest code in the code book.

    Parameters
    ----------
    X : {array-like, sparse matrix}, shape = [n_samples, n_features]
         New data to predict.

    Returns
    -------
    labels : array, shape [n_samples,]
        Index of the cluster each sample belongs to.
    """
    check_is_fitted(self, 'cluster_centers_')

    X = self._check_test_data(X)

    daal_ready = hasattr(X, '__array__') or sp.isspmatrix_csr(X)

    if daal_ready:
        return daal_k_means_dense(X, self.n_clusters, 0, 0.0, self.cluster_centers_, None)[1]
    else:
        x_squared_norms = row_norms(X, squared=True)
        return _labels_inertia(X, x_squared_norms, self.cluster_centers_)[0]
