%matplotlib inline

import sys

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

rng = np.random.RandomState(0)

all_conn_mtx = np.load("data/fmri_autism_conn_mtx.npy")

conn_mtx = np.abs(all_conn_mtx[0, ])

from scipy.sparse.csgraph import minimum_spanning_tree

from dyconnmap.graphs import threshold_omst_global_cost_efficiency

mtx_mst = minimum_spanning_tree(conn_mtx).todense()

nCIJtree, CIJtree, _, _, global_cost_eff_max, cost_max, cost, global_cost_eff = threshold_omst_global_cost_efficiency(conn_mtx)

import matplotlib.pyplot as plt

f, axes = plt.subplots(ncols=3, nrows=1, figsize=(10, 8), dpi=100, sharey=True, sharex=False)

im = axes[0].imshow(conn_mtx, vmin=0.0, vmax=1.0, cmap=plt.cm.Spectral)
axes[0].set_title('Input', fontsize=14)
axes[0].set_ylabel('ROI', fontsize=14)
axes[0].set_xlabel('ROI', fontsize=14)

axes[1].imshow(mtx_mst, vmin=0.0, vmax=1.0, cmap=plt.cm.Spectral)
axes[1].set_title('MST', fontsize=14)
axes[1].set_xlabel('ROI', fontsize=14)

axes[2].imshow(CIJtree, vmin=0.0, vmax=1.0, cmap=plt.cm.Spectral)
axes[2].set_title('OMST', fontsize=14)
axes[2].set_xlabel('ROI', fontsize=14)

f.subplots_adjust(right=0.8)
cbar_ax = f.add_axes([1.0, 0.35, 0.0125, 0.300])
cb = f.colorbar(im, cax=cbar_ax)
cb.set_label('Connectivity', fontsize=14)

plt.tight_layout()
plt.show()

nnz = 32 * 32
nnz_mst = np.count_nonzero(mtx_mst == 0)
nnz_omst = np.count_nonzero(CIJtree == 0)

plt.figure(figsize=(8, 6), dpi=100)
plt.plot([nnz, nnz_mst, nnz_omst], 'o-')
plt.axis('tight')
plt.ylabel('Non-zero elements', fontsize=12)
plt.xticks([0, 1, 2], ['Max', 'MST', 'OMST'], rotation=45, fontsize=12)
plt.grid(alpha=0.25)
plt.title('Sparsity', fontsize=14)
plt.show()