import matplotlib.pyplot as plt
import mnist
import numpy as np

np.random.seed(1234)

plt.rcParams["figure.figsize"] = [16, 9]


def tiles(examples):
    rows_count = examples.shape[0]
    cols_count = examples.shape[1]
    tile_height = examples.shape[2]
    tile_width = examples.shape[3]

    space_between_tiles = 2
    img_matrix = np.empty(shape=( (tile_height + space_between_tiles) * rows_count,(tile_width + space_between_tiles) * cols_count   ))
    img_matrix.fill(np.nan)

    for row in range(rows_count):
        for col in range(cols_count):
            img_matrix[row*(tile_height + space_between_tiles):row*(tile_height + space_between_tiles)+tile_height,   col*(tile_width + space_between_tiles):col*(tile_width + space_between_tiles)+tile_width] = examples[row, col]


    #raise Exception("Not implemented!")

    return img_matrix

def plot_2d_mnist_scatter(X, y):
    fig, plot = plt.subplots()
    fig.set_size_inches(16, 16)
    plt.prism()
    y.shape
    for i in range(10):
        digit_indeces = X[i==y]
        dim1 = digit_indeces[:, 0]
        dim2 = digit_indeces[:, 1]
        plot.scatter(dim1, dim2, label=i)


    plot.set_xticks(())
    plot.set_yticks(())

    plt.tight_layout()
    plt.legend()
    plt.show()



X = mnist.train_images().astype(np.float32) / 255.0
y = mnist.train_labels()

print(X.shape)
X = np.reshape(X,(X.shape[0],-1))
X = X[:10000,:]
y = y[:10000]
print(X.shape)

from sklearn.decomposition import PCA

pca = PCA(n_components=20)
X_pca = pca.fit(X).transform(X)

#print(X_pca.shape)

#plot_2d_mnist_scatter(X_pca,y)

from sklearn.manifold import TSNE

tsne = TSNE(n_iter=300)
x_tsne = TSNE(n_components=2).fit_transform(X)


print(x_tsne.shape)
plot_2d_mnist_scatter(x_tsne,y)