import itertools

import matplotlib.pyplot as plt
from mvlearn.cluster import MultiviewCoRegSpectralClustering, MultiviewKMeans
from mvlearn_adapted.mv_coreg_spectral import MultiviewCoRegSpectralClustering as MVCoReg_adapted

from sklearn.cluster import DBSCAN, MeanShift, KMeans
from sklearn.metrics import normalized_mutual_info_score as nmi_score
import pandas as pd
import numpy as np
from tqdm import tqdm


def adapted_coreg_over_info_views_and_lambda(views, true_labels, lambda_range, n_clusters, n_init=100,
                                     random_state=42, iterate_over_views=False):
    results = {}
    if iterate_over_views:
        views_to_iterate_over = views
    else:
        views_to_iterate_over = [views[0]]
    for most_informative_view in range(len(views_to_iterate_over)):
        print('Most informative view: ', most_informative_view)
        results[most_informative_view] = {}

        pbar = tqdm(lambda_range)

        for lambda_value in pbar:
            pbar.set_description("Processing %s" % lambda_value)
            m_spectral1 = MVCoReg_adapted(n_clusters=n_clusters,
                                                           v_lambda=lambda_value,
                                                           random_state=random_state,
                                                           n_init=n_init,
                                                           affinity='rbf',
                                                           info_view=most_informative_view
                                                           )
            m_clusters1 = m_spectral1.fit_predict(views)
            results[most_informative_view][lambda_value] = nmi_score(true_labels, m_clusters1)
    return results


def coreg_over_info_views_and_lambda(views, true_labels, lambda_range, n_clusters, n_init=100,
                                     random_state=42, iterate_over_views=False):
    results = {}
    if iterate_over_views:
        views_to_iterate_over = views
    else:
        views_to_iterate_over = [views[0]]
    for most_informative_view in range(len(views_to_iterate_over)):
        print('Most informative view: ', most_informative_view)
        results[most_informative_view] = {}

        pbar = tqdm(lambda_range)

        for lambda_value in pbar:
            pbar.set_description("Processing %s" % lambda_value)
            m_spectral1 = MultiviewCoRegSpectralClustering(n_clusters=n_clusters,
                                                           v_lambda=lambda_value,
                                                           random_state=random_state,
                                                           n_init=n_init,
                                                           affinity='rbf',
                                                           info_view=most_informative_view
                                                           )
            m_clusters1 = m_spectral1.fit_predict(views)
            results[most_informative_view][lambda_value] = nmi_score(true_labels, m_clusters1)
    return results


def kmeans_over_random_states(views, true_labels, n_clusters, random_seeds=[0, 1, 42]):
    results = {}

    for random_state in random_seeds:
        m_kmeans = MultiviewKMeans(n_clusters=n_clusters, random_state=random_state)
        m_clusters = m_kmeans.fit_predict(views)
        results[random_state] = nmi_score(true_labels, m_clusters)
    return results


def plot_coreg_over_info_views_and_lambda(results, baselines=None, baseline_labels=None):
    fig = plt.figure(figsize=plt.figaspect(len(results.keys())))

    for info_view in results.keys():
        ax = fig.add_subplot(len(results.keys()), 1, info_view + 1)
        X = []
        Y = []
        for lambda_value in results[info_view].keys():
            X.append(lambda_value)
            Y.append(results[info_view][lambda_value])

        ax.plot(X, Y, marker='.')
        ax.set_ylim(0, 1)
        ax.set_title('Most informative view: ' + str(info_view + 1))
        ax.set_xlabel('Lambda')
        ax.set_ylabel('NMI')

        if baselines is not None:
            color = plt.cm.spring(np.linspace(0, 1, len(baselines)))
            for i, baseline in enumerate(baselines):
                ax.axhline(baseline, label=baseline_labels[i], ls='--', c=color[i])
        ax.legend()
    return fig


def get_view_combinations(views):
    all_view_combinations = []
    view_indices = []
    for i in range(2, len(views) + 1):
        for j, combo in enumerate(itertools.combinations(views, i)):
            all_view_combinations.append(combo)
            view_indices.append(list(itertools.combinations(range(len(views)), i))[j])
    return all_view_combinations, view_indices


def plot_combination_results(combination_results, view_indices):
    combination_lengths = len(set([len(res.keys()) for res in combination_results]))

    number_of_views_iterator = len(combination_results[0].keys())
    j = 0

    fig, axs = plt.subplots(combination_lengths, figsize=(15, 5 * combination_lengths))

    for i, results in enumerate(combination_results):
        # colors = cm.rainbow(np.linspace(0, 1, len(results.keys())))

        number_of_views = len(results.keys())
        if number_of_views != number_of_views_iterator:
            number_of_views_iterator = number_of_views
            j = j + 1

        for info_view in results.keys():
            # ax = fig.add_subplot(len(results.keys()), 1, info_view + 1)
            X = []
            Y = []
            for lambda_value in results[info_view].keys():
                X.append(lambda_value)
                Y.append(results[info_view][lambda_value])
            axs[j].plot(X, Y, marker='.',
                        label='views ' + str(view_indices[i]) + 'miv:' + str(info_view))
        axs[j].set_title('Number of views:' + str(number_of_views))
        axs[j].set_ylim(0, 1)
        axs[j].set_xlabel('Lambda')
        axs[j].set_ylabel('NMI')
        axs[j].legend(loc='center left', bbox_to_anchor=(1, 0.5), ncol=3)
    fig.tight_layout()


def print_combination_results(combination_results, view_indices):
    pd.options.display.float_format = '{:,.4f}'.format

    number_of_views_iterator = len(combination_results[0].keys())
    j = 0

    for i, results in enumerate(combination_results):
        print(str(i) + '--------------------------------------------------')
        # df = pd.Dataframe(columns=['Combination',  'Views', 'Lambda', 'NMI'])

        number_of_views = len(results.keys())
        if number_of_views != number_of_views_iterator:
            number_of_views_iterator = number_of_views
            j = j + 1

        for info_view in results.keys():

            # ax = fig.add_subplot(len(results.keys()), 1, info_view + 1)
            X = []
            Y = []

            for lambda_value in results[info_view].keys():
                X.append(lambda_value)
                Y.append(results[info_view][lambda_value])

            df = pd.DataFrame(
                {
                    "Combination": str(view_indices[i]),
                    "Info_view": str(info_view),
                    "lambda": X,
                    "NMI": Y
                }
            )

            print(df)


def print_k_means_results(combination_results, view_indices):
    pd.options.display.float_format = '{:,.4f}'.format

    number_of_views_iterator = len(combination_results[0].keys())
    j = 0

    for i, results in enumerate(combination_results):
        print(str(i) + '--------------------------------------------------')
        # df = pd.Dataframe(columns=['Combination',  'Views', 'Lambda', 'NMI'])

        number_of_views = len(results.keys())
        if number_of_views != number_of_views_iterator:
            number_of_views_iterator = number_of_views
            j = j + 1

        X = []
        Y = []

        for random_state in results.keys():
            # ax = fig.add_subplot(len(results.keys()), 1, info_view + 1)

            X.append(random_state)
            Y.append(results[random_state])

        df = pd.DataFrame(
            {
                "Combination": str(view_indices[i]),
                "Random_state": X,
                "NMI": Y
            }
        )

        print(df)

def run_k_means(views, labels, num_of_clusters):
    s_kmeans = KMeans(
        n_clusters=num_of_clusters, random_state=1, init='k-means++')

    for i in range(len(views)):
        s_clusters_kmeans = s_kmeans.fit_predict(views[i])
        s_nmi_kmeans = nmi_score(labels, s_clusters_kmeans, average_method='arithmetic')
        print('Single-view View {0:d} NMI Score k-means: {1:.3f}\n'.format(i, s_nmi_kmeans))

    s_data = np.hstack(views)

    s_clusters_kmeans = s_kmeans.fit_predict(s_data)
    s_nmi_kmeans = nmi_score(labels, s_clusters_kmeans)
    print('Single-view Concatenated NMI Score k-means: {0:.3f}\n'.format(s_nmi_kmeans))

def run_dbscan_for_eps(views, labels, eps):
    s_dbscan = DBSCAN(eps=eps)

    for i in range(len(views)):
        s_cluster_dbscan = s_dbscan.fit_predict(views[i])
        s_nmi_dbscan = nmi_score(labels, s_cluster_dbscan, average_method='arithmetic')

        print('Single-view View {0:d} NMI Score DBSCAN eps={1:.3f}: {2:.3f}\n'
              .format(i, eps, s_nmi_dbscan))

    s_data = np.hstack(views)

    s_clusters_dbscan = s_dbscan.fit_predict(s_data)
    s_nmi_dbscan = nmi_score(labels, s_clusters_dbscan)
    print('Single-view Concatenated NMI Score DBSCAN: {0:.3f}\n'.format(s_nmi_dbscan))
    print('-------------------------------------------------------------------')


def run_meanshift_for_bandwidth(views, labels, bandwidth):
    s_meanshift = MeanShift(bandwidth=bandwidth)

    for i in range(len(views)):
        s_cluster_meanshift = s_meanshift.fit_predict(views[i])
        s_nmi_meanshift = nmi_score(labels, s_cluster_meanshift, average_method='arithmetic')

        print('Single-view View {0:d} NMI Score Meanshift bandwidth={1:.3f}: {2:.3f}\n'
              .format(i, bandwidth, s_nmi_meanshift))

    s_data = np.hstack(views)

    s_cluster_meanshift = s_meanshift.fit_predict(s_data)
    s_nmi_meanshift = nmi_score(labels, s_cluster_meanshift)
    print('Single-view Concatenated NMI Score Meanshift: {0:.3f}\n'.format(s_nmi_meanshift))
    print('-------------------------------------------------------------------')

# view_1 = [[0,2],[0,3],[2,0]]
# view_2 = [[1,2], [1,3], [3,1]]
# view_3 = [[1,2],[1,3],[3,0]]
#
# kmeans_over_random_states(views=[view_1, view_2, view_3], true_labels=[0,0,1], n_clusters=2, random_seeds=[0,1])
