%matplotlib inline

import tensorflow as tf
from tensorflow.keras import layers
import numpy as np
from scipy.cluster import hierarchy
import fastcluster
from statsmodels.stats.correlation_tools import corr_nearest
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings("ignore")


# dim in [75, 80, 100, 120, 140, 160, 180 200]
dim = 80
nsamples = 4

generator = tf.keras.models.load_model(
    'saved_model/generator_{}d'.format(dim))

noise_dim = 100
noise = tf.random.normal([nsamples, noise_dim])
generated_image = generator(noise, training=False)
    
a, b = np.triu_indices(dim, k=1)

nearest_corrmats = []
for i in range(nsamples):
    corrmat = np.array(generated_image[i, :, :, 0])
    # set diag to 1
    np.fill_diagonal(corrmat, 1)
    # symmetrize
    corrmat[b, a] = corrmat[a, b]
    # nearest corr
    nearest_corrmat = corr_nearest(corrmat)
    # set diag to 1
    np.fill_diagonal(nearest_corrmat, 1)
    # symmetrize
    nearest_corrmat[b, a] = nearest_corrmat[a, b]

    # arrange with hierarchical clustering
    dist = 1 - nearest_corrmat
    dim = len(dist)
    tri_a, tri_b = np.triu_indices(dim, k=1)
    Z = fastcluster.linkage(dist[tri_a, tri_b], method='ward')
    permutation = hierarchy.leaves_list(
    hierarchy.optimal_leaf_ordering(Z, dist[tri_a, tri_b]))
    ordered_corr = nearest_corrmat[permutation, :][:, permutation]

    nearest_corrmats.append(ordered_corr)

plt.figure(figsize=(12, 8))
for i in range(min(4, len(nearest_corrmats))):
    plt.subplot(2, 2, i + 1)
    plt.pcolormesh(nearest_corrmats[i][:, :], cmap='viridis')
    plt.colorbar()
plt.show()