# embedding
from sklearn.manifold import TSNE
import pandas as pd
import matplotlib.pyplot as plt
from plotnine import aes, geom_point, ggplot, theme_bw


def visualize_embedding(multinet, labels=None, verbose=True):
    embedding = multinet.embedding
    X = embedding[0]
    indices = embedding[1]

    if verbose:
        print("------ Starting embedding visualization -------")

    if labels:
        # optionally match indices to labels and add a column
        label_vector = [labels[x] for x in indices]
        X_embedded = TSNE(n_components=2).fit_transform(X)
        dfr = pd.DataFrame(X_embedded, columns=['dim1', 'dim2'])
        dfr['labels'] = label_vector
        print(dfr.head())
        gx = (ggplot(dfr, aes('dim1', 'dim2', color="labels")) +
              geom_point(size=0.5) + theme_bw())
        gx.draw()
        plt.show()
    else:
        X_embedded = TSNE(n_components=2).fit_transform(X)
        dfr = pd.DataFrame(X_embedded, columns=['dim1', 'dim2'])
        print(dfr.head())
        gx = (ggplot(dfr, aes('dim1', 'dim2')) + geom_point(size=0.5) +
              theme_bw())
        gx.draw()
        plt.show()
