import sys

import numpy as np
np.set_printoptions(threshold=sys.maxsize)

labels = np.load('data/fmri_autism_labels.npy')
data = np.load('data/fmri_autism_thres_conn_mtx.npy')

labels = np.array(labels).astype(np.int32)

con_indices = np.where(labels == 0)[0]
aut_indices = np.where(labels == 1)[0]

con_indices = con_indices[0:10]
aut_indices = aut_indices[0:10]

ordered_data = np.vstack([data[con_indices, :, :],
                          data[aut_indices, :, :],
])

ordered_data = np.abs(ordered_data)

num_subjects, num_rois, num_rois = np.shape(ordered_data)

print(f"""
            Subjects: {num_subjects}
Coonnectivity Matrix: {num_rois}x{num_rois}
""")

"""
Explanation: In this short tutorial, we will demonstrate how to use measure the distance between two graphs, based on the supplied methods.

Load and prepare data
End of explanation
"""


from dyconnmap.graphs import graph_diffusion_distance, im_distance, spectral_euclidean_distance, spectral_k_distance

methods = {
    'gdd': graph_diffusion_distance,
    'imd': im_distance,
    'spectral': spectral_euclidean_distance,
    'k_spectral': spectral_k_distance
}

results = dict()

for name, func in methods.items():
    print('- Running method {0}...'.format(name), end=' ')
    
    distances = np.zeros((num_subjects, num_subjects))
    
    for x in range(num_subjects):
        for y in range(num_subjects):            
            mtx1 = ordered_data[x, :, :]
            mtx2 = ordered_data[y, :, :]
            
            # Workdaround until `spectral_k_distance`
            # gets a default value in its implementation.
            if name == 'k_spectral':
                method_distance = func(mtx1, mtx2, k=5)
            else:
                method_distance = func(mtx1, mtx2)
            
            if name == 'gdd':
                method_distance = method_distance[0]
            
            distances[x, y] = method_distance
    
    results[name] = distances
    
    print('Done.')

"""
Explanation: Estimate the distances
End of explanation
"""


import matplotlib.pyplot as plt

plot_labels = np.hstack([10 * ['CON'], 10 * ['AUT']])

f, axes = plt.subplots(ncols=4, figsize=(16, 4), sharex=True, sharey=True)

for i, (name, mtx) in enumerate(results.items()):
    im = axes[i].imshow(mtx, cmap=plt.cm.Spectral)
    f.colorbar(im, ax=axes[i], )
    axes[i].set_title(name)
    axes[i].set_aspect('auto')
    axes[i].set_xticks(range(20))
    axes[i].set_xticklabels(plot_labels, fontsize=8, rotation=90)

plt.yticks(range(20), plot_labels, fontsize=8)
plt.tight_layout(h_pad=1)
plt.show()

"""
Explanation: Plot
Plot the distance matrices
End of explanation
"""


from sklearn.manifold import MDS

mds = MDS(n_components=2, dissimilarity='precomputed', random_state=0)
xy = mds.fit_transform(results['spectral'])

plot_labels = np.vstack([np.ones((10, 1)), np.ones((10, 1)) * 2]).squeeze()

plt.figure(figsize=(6, 6))
s1 = plt.scatter(xy[0:10, 0], xy[0:10, 1], marker='x', c='g', s=50, edgecolors='w', zorder=1000, label='CON')
s2 = plt.scatter(xy[10:, 0], xy[10:, 1], marker='d', c='b', s=50, edgecolors='w', zorder=1000, label='AUT')
plt.axis('equal')
plt.grid(alpha=0.25)
plt.legend()
plt.title('MDS GDD')
plt.show()

"""
Explanation: Project the Graph Diffusion Distance matrix to two dimensions, using MDS
End of explanation
"""


