import tensorflow as tf

import numpy as np

from glob import glob

from natsort import natsorted

import os

from model import TripleNet, train_step, test_step

from utils import load_complete_data

from tqdm import tqdm

from sklearn.manifold import TSNE

import matplotlib.pyplot as plt

from matplotlib import style

import seaborn as sns

import pandas as pd

import pickle

from sklearn.cluster import KMeans

from scipy.optimize import linear_sum_assignment as linear_assignment

style.use('seaborn')

os.environ["CUDA_DEVICE_ORDER"]= "PCI_BUS_ID"

os.environ["CUDA_VISIBLE_DEVICES"]= '0'

# Thanks to: https://github.com/k-han/DTC/blob/master/utils/util.py

def cluster_acc(y_true, y_pred):

    """

    Calculate clustering accuracy. Require scikit-learn installed

    # Arguments

        y: true labels, numpy.array with shape `(n_samples,)`

        y_pred: predicted labels, numpy.array with shape `(n_samples,)`

    # Return

        accuracy, in [0,1]

    """

    y_true = y_true.astype(np.int64)

    assert y_pred.size == y_true.size

    D = max(y_pred.max(), y_true.max()) + 1

    w = np.zeros((D, D), dtype=np.int64)

    for i in range(y_pred.size):

        w[y_pred[i], y_true[i]] += 1

    ind = linear_assignment(w.max() - w)

    return sum([w[i, j] for i, j in zip(*ind)]) * 1.0 / y_pred.size

if __name__ == '__main__':

    n_channels  = 14

    n_feat      = 128

    batch_size  = 256

    test_batch_size  = 256

    n_classes   = 10

    # data_cls = natsorted(glob('data/thoughtviz_eeg_data/*'))

    # cls2idx  = {key.split(os.path.sep)[-1]:idx for idx, key in enumerate(data_cls, start=0)}

    # idx2cls  = {value:key for key, value in cls2idx.items()}

    with open('../../data/b2i_data/eeg/image/data.pkl', 'rb') as file:

        data = pickle.load(file, encoding='latin1')

        train_X = data['x_train']

        train_Y = data['y_train']

        test_X = data['x_test']

        test_Y = data['y_test']

    # train_batch = load_complete_data('data/thoughtviz_eeg_data/*/train/*', batch_size=batch_size)

    # val_batch   = load_complete_data('data/thoughtviz_eeg_data/*/val/*', batch_size=batch_size)

    # test_batch  = load_complete_data('data/thoughtviz_eeg_data/*/test/*', batch_size=test_batch_size)

    train_batch = load_complete_data(train_X, train_Y, batch_size=batch_size)

    val_batch   = load_complete_data(test_X, test_Y, batch_size=batch_size)

    test_batch  = load_complete_data(test_X, test_Y, batch_size=test_batch_size)

    # X, Y = next(iter(train_batch))

    # print(X.shape, Y.shape)

    triplenet = TripleNet(n_classes=n_classes)

    opt     = tf.keras.optimizers.Adam(learning_rate=3e-4)

    triplenet_ckpt    = tf.train.Checkpoint(step=tf.Variable(1), model=triplenet, optimizer=opt)

    triplenet_ckpt.restore('experiments/best_ckpt/ckpt-89')

    tq = tqdm(test_batch)

    feat_X  = np.array([])

    feat_Y  = np.array([])

    for idx, (X, Y) in enumerate(tq, start=1):

        _, feat = triplenet(X, training=False)

        feat_X = np.concatenate((feat_X, feat.numpy()), axis=0) if feat_X.size else feat.numpy()

        feat_Y = np.concatenate((feat_Y, Y.numpy()), axis=0) if feat_Y.size else Y.numpy()

    print(feat_X.shape, feat_Y.shape)

    # colors = list(plt.cm.get_cmap('viridis', 10))

    # print(colors)

    # colors  = [np.random.rand(3,) for _ in range(10)]

    # print(colors)

    # Y_color = [colors[label] for label in feat_Y]

    kmeans = KMeans(n_clusters=n_classes,random_state=45)

    kmeans.fit(feat_X)

    labels = kmeans.labels_

    kmeanacc = cluster_acc(feat_Y, labels)

    # correct_labels = sum(feat_Y == labels)

    # print("Result: %d out of %d samples were correctly labeled." % (correct_labels, feat_Y.shape[0]))

    # kmeanacc = correct_labels/float(feat_Y.shape[0])

    print('Accuracy score: {0:0.2f}'. format(kmeanacc))

    tsne = TSNE(n_components=2, verbose=1, perplexity=40, n_iter=700)

    tsne_results = tsne.fit_transform(feat_X)

    df = pd.DataFrame()

    df['label'] = feat_Y

    df['x1'] = tsne_results[:, 0]

    df['x2'] = tsne_results[:, 1]

    # df['x3'] = tsne_results[:, 2]

    df.to_csv('experiments/infer_triplet_embed2D.csv')  

    # df.to_csv('experiments/triplenet_embed3D.csv')

    # df = pd.read_csv('experiments/triplenet_embed2D.csv')

    df = pd.read_csv('experiments/infer_triplet_embed2D.csv')

    plt.figure(figsize=(16,10))

    # ax = plt.axes(projection='3d')

    sns.scatterplot(

        x="x1", y="x2",

        data=df,

        hue='label',

        palette=sns.color_palette("hls", n_classes),

        legend="full",

        alpha=0.4

    )

    # ax.scatter3D(df['x1'], df['x2'], df['x3'], c=df['x3'], alpha=0.4)

    # plt.scatter(df['x1'], df['x2'], c=df['x2'], alpha=0.4)

    # min_x, max_x = np.min(feat_X), np.max(feat_X)

    # min_x, max_x = -10, 10

    # for c in range(len(np.unique(feat_Y))):

    #   # ax.scatter(feat_X[feat_Y==c, 0], feat_X[feat_Y==c, 1], feat_X[feat_Y==c, 2], '.', alpha=0.5, c=colors[c], s=0.3)

    #   plt.scatter(feat_X[feat_Y==c, 0], feat_X[feat_Y==c, 1], marker='.', alpha=0.5, c=colors[c], s=1.0)

    # plt.title('Triple Loss')

    # W = triplenet.cls_layer.get_weights()[0].T

    # x = np.linspace(min_x, max_x, 50)

    # y = W[0][1]*x + W[0][0]

    # plt.plot(x, y, c=colors[0])

    # x = np.linspace(min_x, max_x, 50)

    # y = W[1][1]*x + W[1][0]

    # plt.plot(x, y, c=colors[1])

    # x = np.linspace(min_x, max_x, 50)

    # y = W[2][1]*x + W[2][0]

    # plt.plot(x, y, c=colors[2])

    # x = np.linspace(min_x, max_x, 50)

    # y = W[3][1]*x + W[3][0]

    # plt.plot(x, y, c=colors[3])

    # x = np.linspace(min_x, max_x, 50)

    # y = W[4][1]*x + W[4][0]

    # plt.plot(x, y, c=colors[4])

    # plt.clf()

    # plt.close()

    # featX = df[['x1', 'x2']].to_numpy()

    # print(featX.shape)

    plt.title('k-means accuracy: {}%'.format(kmeanacc*100))

    plt.savefig('experiments/embedding.png')

    # plt.show()