# Author: Andrey Boytsov <andrey.boytsov@uni.lu> <andrey.m.boytsov@gmail.com>
# License: BSD 3 clause (C) 2017

# Visualizing 8x8 digits dataset (from sklearn). Visualizing first half, then generating embedding function and
# embedding the rest.

import sys
import os
from sklearn.datasets import load_digits
import matplotlib.pyplot as plt

# Importing from parent directory
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
import dynamic_tsne

if __name__ == "__main__":
    data = load_digits()
    X_all = data.images.reshape((-1, 64)) # 8x8 image to 64-length vector
    X = X_all[:len(X_all)//2, :]
    X2 = X_all[len(X_all)//2:, :]
    labels_all = data.target
    labels = labels_all[:len(labels_all)//2]
    labels2 = labels_all[len(labels_all)//2:]

    dTSNE = dynamic_tsne.DynamicTSNE(perplexity=20)
    y = dTSNE.fit(X, verbose=2, optimizer_kwargs={'momentum': 0.8}, random_seed=0)
    embedder = dTSNE.generate_embedding_function(embedding_function_type='weighted-inverse-distance',
                                                 function_kwargs={'power': 20})
    y2 = embedder(X2, verbose=2)

    color_list = ['blue', 'red', 'green', 'yellow', 'cyan', 'black', 'magenta', 'pink', 'brown', 'orange']

    plt.gcf().set_size_inches(10, 10)
    legend_list = list()
    for l in set(sorted(labels)):
        plt.scatter(y[labels == l, 0], y[labels == l, 1], c=color_list[l])
        legend_list.append(str(data.target_names[l]))
    for l in set(sorted(labels)):
        plt.scatter(y2[labels2 == l, 0], y2[labels2 == l, 1], c=color_list[l], marker='v')
        legend_list.append(str(data.target_names[l])+" embedded")
    # It did not work. Clusters were elsewhere.
    #plt.xlim([-70,90])
    #plt.ylim([-90,70])
    plt.legend(legend_list)
    plt.show()
