import logging

import dask_ml.cluster
import hdbscan
import pandas as pd
from sklearn import metrics
from sklearn.cluster import KMeans, MiniBatchKMeans

from ..plot_metrics import draw_graphs

logging.getLogger().setLevel(logging.INFO)

# We set dask_ml logging to ERROR, otherwise our logs get bloated by crap
dask_logger = logging.getLogger("dask_ml.cluster.k_means")
dask_logger.setLevel(logging.ERROR)


def run_algo(training_set, fan_ids, modelling_config: dict, algo_params: dict):

    X = training_set
    if isinstance(fan_ids, pd.Series):
        fan_ids = fan_ids.to_frame()

    hyperparams = algo_params['hyperparams']
    n_clusters = algo_params['n_clusters']
    algorithm = algo_params['algorithm']

    hyperparams['n_clusters'] = n_clusters
    hyperparams['random_state'] = 0  # For consistency over multiple runs, we want this to always remain the same

    if algorithm == 'sklearn':
        kmeans = KMeans(**hyperparams).fit(X)
        labels = kmeans.labels_
    elif algorithm == 'sklearn_minibatch':
        kmeans = MiniBatchKMeans(**hyperparams, batch_size=5000).fit(X)
        labels = kmeans.labels_
    else:  # algorithm == 'dask'
        hyperparams['n_jobs'] = -1  # We always want to use all CPU cores
        kmeans = dask_ml.cluster.KMeans(**hyperparams).fit(X)
        labels = kmeans.labels_.compute()  # from dask_ml, we get dask array, which we convert to numpy array

    centers = kmeans.cluster_centers_
    # TODO maybe remove metrics that create NxN distance matrix
    validity_index = hdbscan.validity.validity_index(X, labels, metric="euclidean") if algorithm == 'dask' else 1

    cluster_output = fan_ids.join(pd.DataFrame(labels)).rename(columns={0: 'cluster', '0': 'cluster'})
    X = pd.DataFrame(X).join(pd.DataFrame(labels), lsuffix='left_').rename(columns={0: 'cluster', '0': 'cluster'})
    # TODO optimize drawing plots for big datasets or not draw them on Live at all
    graph_locations = draw_graphs(X, labels, modelling_config, n_clusters, centers)

    eval_metrics = {'clusters': n_clusters,
                    'intertia': kmeans.inertia_,
                    'n_iter': kmeans.n_iter_,
                    'silhouette_coefficient': metrics.silhouette_score(X, labels),
                    'davies_bouldin_index': metrics.davies_bouldin_score(X, labels),
                    'calinski_harabasz_index': metrics.calinski_harabasz_score(X, labels),
                    'validity_index': validity_index
                    }

    return kmeans, eval_metrics, hyperparams, cluster_output, graph_locations


if __name__ == '__main__':
    """ Instead of generating df, perhaps load your own dataframe from csv. We need a dataframe with
    training cols only (preprocessed, and standardscaled) and a fan_id. """
    data = pd.DataFrame({
        'fan_id': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14],
        'a': [1,2,3,4,5,6,4,3,1,5,5,6,7,9],
        'b': [2,3,4,5,6,7,8,9,2,3,4,5,6,9]
    })
    training_set = data[['a', 'b']]
    fan_ids = data['fan_id']

    modelling_config = {
        'plot_2d': True,
        'plot_3d': True,
        'plot_silhouette': True,
        'boxplot': True,
    }
    algo_params = {'n_clusters': 4, 'hyperparams': {'init': 'k-means++', 'max_iter': 500}}

    print(run_algo(training_set, fan_ids, modelling_config, algo_params)[1])