import numpy as np
from matplotlib import gridspec
from matplotlib import pyplot as plt

from .basic import normalize

_fontsize = 14


def plot_point_process(arrivals, intensities, true_intensity=None):
    plt.figure(figsize=(10, 5))
    gs = gridspec.GridSpec(2, 1, height_ratios=[4, 1])
    ax0 = plt.subplot(gs[0])
    ax0.set_ylabel(r'Hawkes Intensity $\lambda(t)$', fontsize=20)
    ax0.plot(np.sort(arrivals), intensities)

    ax1 = plt.subplot(gs[1])
    ax1.spines['right'].set_visible(False)
    ax1.spines['left'].set_visible(False)
    ax1.get_yaxis().set_visible(False)
    ax1.vlines(arrivals, ymin=0, ymax=1,
               color="b", label="arrivals")

    if true_intensity is not None:
        ax0.plot(true_intensity[0], true_intensity[1], "k-",
                 color="r", label="sine function", linewidth=3)
    plt.xlabel(r'Time ($t$)', fontsize=20)
    plt.tick_params(axis='both', which='major', labelsize=14)
    plt.tick_params(axis='both', which='minor', labelsize=14)
    ax0.set_xlim(np.max(arrivals)*.003, np.max(arrivals))
    ax1.set_xlim(np.max(arrivals)*.003, np.max(arrivals))
    plt.pause(500)


def plot_intensities(arrivals, intensities, support, true_intensity):
    fig = plt.figure(figsize=(10, 5))
    gs = gridspec.GridSpec(2, 1, height_ratios=[4, 1])
    ax0 = plt.subplot(gs[0])
    ax0.set_ylabel(r'Hawkes Intensity $\lambda(t)$', fontsize=20)
    #       ax0.get_xaxis().set_visible(False)
    #       ax0.tick_params(labelleft='off')
    plt.plot(intensities[0], intensities[1])

    ax1 = plt.subplot(gs[1])
    #         ax1.tick_params(labelbottom='off')
    ax1.spines['right'].set_visible(False)
    ax1.spines['left'].set_visible(False)
    ax1.get_yaxis().set_visible(False)

    lns1 = ax0.plot(support, true_intensity, "k-", color="r",
                    label="sine function", linewidth=3)

    lns3 = ax1.vlines(arrivals, ymin=0, ymax=1, color="b", label="arrivals")

    plt.xlabel(r'Time ($t$)', fontsize=20)

    plt.tick_params(axis='both', which='major', labelsize=14)
    plt.tick_params(axis='both', which='minor', labelsize=14)


def plot_final_intensities(f, model, sess, input_data):
    intensities = model.inference(input_data, sess)
    x = input_data[0, :, :, 0]
    n, m = x.shape
    x = np.sort(x.reshape((n * m)))

    fig = plt.figure(figsize=(10, 5))
    gs = gridspec.GridSpec(2, 1, height_ratios=[4, 1])
    ax0 = plt.subplot(gs[0])
    ax0.set_ylabel(r'Hawkes Intensity $\lambda(t)$', fontsize=20)
    #       ax0.get_xaxis().set_visible(False)
    #       ax0.tick_params(labelleft='off')

    plt.plot(x, intensities)

    max_arr = np.max(x)
    support = np.arange(0, max_arr, .01)
    ax1 = plt.subplot(gs[1])
    #         ax1.tick_params(labelbottom='off')
    ax1.spines['right'].set_visible(False)
    ax1.spines['left'].set_visible(False)
    ax1.get_yaxis().set_visible(False)
    lns1 = ax0.plot(support, f(support), "k-", color="r",
                    label="sine function", linewidth=3)
    lns3 = ax1.vlines(x, ymin=0, ymax=1,
                      color="b", label="arrivals")
    plt.xlabel(r'Time ($t$)', fontsize=20)
    plt.tick_params(axis='both', which='major', labelsize=14)
    plt.tick_params(axis='both', which='minor', labelsize=14)

    plt.show()
    plt.pause(500.05)


def find_best_model_using(clustering_models, measure='silhouette_score'):
    score = []
    clusters = []
    for c in clustering_models:
        clusters.append(c)
        score.append(clustering_models[c][measure])
    return clustering_models[clusters[np.argmax(score)]], clusters[np.argmax(score)]


def plot_all_models_performance(models, title=None, measure='silhouette_score', save_to_file=None):
    for model_name in models:
        model = models[model_name]['model']
        color = models[model_name]['color']
        label = models[model_name]['label']
        sizes = []
        silhouettes = []
        inertias = []
        first = True
        for clustering_size in model:
            silhouette_score = model[clustering_size]['silhouette_score']
            inertia = model[clustering_size]['inertia']
            sizes.append(clustering_size)
            silhouettes.append(silhouette_score)
            inertias.append(inertia)
        indexs_to_order_by = np.array(sizes, dtype=np.int32).argsort()
        sizes_ordered = np.array(sizes)[indexs_to_order_by]
        silhouettes_ordered = np.array(silhouettes)[indexs_to_order_by]
        inertias_ordered = np.array(inertias)[indexs_to_order_by]
        plt.grid(True)
        if not title is None:
            plt.title(title, fontsize=_fontsize)
        plt.gca().spines['top'].set_visible(False)
        plt.gca().spines['right'].set_visible(False)
        plt.xlabel('Number of Clusters', fontsize=_fontsize)
        if measure == 'silhouette_score':
            plt.ylabel('Average Silhouette Coefficient', fontsize=_fontsize)
            plt.plot(sizes_ordered, silhouettes_ordered,
                     color=color, label=label)
        else:
            plt.ylabel('Inertia', fontsize=_fontsize)
            plt.plot(sizes_ordered, inertias_ordered, color=color, label=label)

    if not save_to_file is None:
        plt.savefig(save_to_file)
        plt.close()
    plt.show()


def plot_clustering(data, clustering_model, is_normalize=False, min_x=1, max_x=365., save_to_file=None):
    total = len(data)

    for j, centroid in enumerate(clustering_model['centroids']):
        cluster_size = clustering_model['cluster_size'][j]
        plt.title('{0:.2%}'.format(cluster_size / float(total)), fontsize=26)
        for i in range(cluster_size):
            plt.plot(np.arange(
                min_x, max_x), data[clustering_model['labels'] == j][i], alpha=.03, color='#168EF7')
        if is_normalize:
            centroid = normalize(0., 1., centroid)
        # plt.tight_layout()
        plt.gcf().subplots_adjust(bottom=0.15)
        plt.xlim(min_x, max_x)
        plt.xticks([min_x, (min_x + max_x) / 2., max_x], fontsize=14)
        plt.xlabel('Day', fontsize=20)
        plt.ylabel('Intensity', fontsize=20)
        plt.yticks([0, 0.5, 1], fontsize=14)
        plt.gca().spines['top'].set_visible(False)
        plt.gca().spines['right'].set_visible(False)
        plt.plot(np.arange(0, len(centroid)), centroid,
                 linewidth=3.0, color='#F88017')

        if not save_to_file is None:
            plt.savefig("{1}_cluster_{0}.pdf".format(j, save_to_file))
            plt.close()
        else:
            plt.show()


def find_best_model_from_dictionary_of_models(models, method='silhouette_score'):
    best_model = {'silhouette_score': -100, 'inertia': 1e10}
    best_model_name = ''
    for model_name in models:
        model = models[model_name]['model']

        current_model, ass = find_best_model_using(model, method)
        if method == 'silhouette_score':
            if current_model[method] > best_model[method]:
                best_model = current_model
                best_model_name = model_name
        else:
            if current_model[method] < best_model[method]:
                best_model = current_model
                best_model_name = model_name

    return best_model


def find_best_model_give_number_of_clusters(models, number_of_clusters):
    silhouette_scores = []
    model_names = []
    for model in models:
        for c in models[model]['model']:
            if c == number_of_clusters:
                model_names.append(model)
                silhouette_scores.append(
                    models[model]['model'][c]['silhouette_score'])
    best_model_name = model_names[np.argmax(silhouette_scores)]
    return models[best_model_name]['model'][number_of_clusters], best_model_name


def plot_inferred_user_intesity(complete_model, packed_model, save_to_file=None):
    x_t = np.arange(complete_model.h.min_arrival,
                    complete_model.h.max_arrival, 0.1)
    fig = plt.figure(figsize=(12, 6))
    gs = gridspec.GridSpec(3, 1, height_ratios=[7, 1, 1])
    ax0 = plt.subplot(gs[0])
    ax0.plot(x_t, complete_model.mu(x_t), "k-",
             label="Complete Model", color="#168EF7")
    ax0.plot(x_t, packed_model.mu(x_t), "k-",
             label="Packed Model", color="#F88017")
    ax0.set_ylabel(r'Poisson Intensity $\lambda(t)$', fontsize=18)

    ax0.get_xaxis().set_visible(False)
    ax0.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=3,
               ncol=2, mode="expand", borderaxespad=0.)
    ax1 = plt.subplot(gs[1])

    ax1.vlines(complete_model.h.X, ymin=0, ymax=1,
               color="#168EF7", label="All Arrivals")

    ax2 = plt.subplot(gs[2])
    ax2.vlines(packed_model.h.X, ymin=0, ymax=1, color="#F88017",
               alpha=.9, label="Packed Arrivals")

    plot_limit = (complete_model.h.min_arrival, complete_model.h.max_arrival)
    ax0.set_xlim(plot_limit)
    ax1.set_xlim(plot_limit)
    ax1.get_yaxis().set_visible(False)
    ax1.get_xaxis().set_visible(False)
    ax2.set_xlim(plot_limit)
    ax2.get_yaxis().set_visible(False)
    ax2.set_xlabel(r'Time ($t$)', fontsize=18)
    # fig.tight_layout()

    ax0.tick_params(axis='both', which='major', labelsize=14)
    ax1.tick_params(axis='both', which='minor', labelsize=14)
    ax2.tick_params(axis='both', which='minor', labelsize=14)
    if save_to_file is None:
        plt.show()
    else:
        plt.savefig(save_to_file)
    plt.close()


def plot_compressed_arrivals(packages, compressed, arrivals):
    fig = plt.figure(figsize=(12, 3))
    gs = gridspec.GridSpec(2, 1, height_ratios=[1, 1])

    ax0 = plt.subplot(gs[0])
    ax0.vlines(compressed, ymin=0, ymax=1, color="r")

    ax1 = plt.subplot(gs[1])
    ax1.vlines(arrivals, ymin=0, ymax=1, color="b")

    plot_limit = (np.min(arrivals), np.max(arrivals))

    ax0.set_xlim(plot_limit)
    ax0.get_yaxis().set_visible(False)
    ax1.set_xlim(plot_limit)
    ax1.get_yaxis().set_visible(False)
    for l in packages:
        if len(l) > 0:
            ax1.fill_between(l, 0, 1, facecolor='green', alpha=.5, step='mid')

    ax1.set_xlabel(r'Time ($t$)', fontsize=18)
    fig.tight_layout()

    ax1.tick_params(axis='both', which='minor', labelsize=14)
