text: '''
  This modules contains functions necessary for retraining the model after the
  initial training with OWL or group OWL
  '''
  from __future__ import division, print_function, absolute_import
  import sys
  sys.path.append('./growl_param_share') #TODO
  sys.path.append('./owl_param_share') #TODO
  
  import tensorflow as tf
  import numpy as np
  from sklearn.cluster import AffinityPropagation
  import utils_nn
  from numpy.linalg import norm
  from sklearn.preprocessing import normalize
  from utils_plot import imagesc
  import matplotlib.pyplot as plt
  import yaml
  
  
  def display_similarity(sess, epoch, get_group, config):
  
      '''
      display the pairwise similarity of the rows of the weight matrix. and plot the
      clustering results.
  
      Args:
          sess: the computation graph
          epoch: current epoch
  
      Return:
          num_nonzero_rows_tuple: # of nonzero rows for each layer
          num_clusters_tuple: # of clusters for each layer
          group_info: list of tuples containing the rows indices of each group
  
      '''
  
      num_nonzero_rows_tuple, num_clusters_tuple = tuple(), tuple()
  
      mask_placeholders = utils_nn.get_mask_placeholders()
      weight_placeholders = utils_nn.get_weight_placeholders()
  
      num_layers = 0
  
      threshold = np.finfo(np.float32).eps
  
      assert len(mask_placeholders) == len(weight_placeholders)
      group_info = []
  
      # only process the layers we applied grOWL on
      for idx,  mask_triple in enumerate(mask_placeholders):
  
          # double check we have turned on grOWL for the layer
          if not config['owl_applied_layers'][idx]:
              continue
  
          num_layers = num_layers + 1
          mask_i, mask_placeholders_i, mask_assign_op_i = mask_placeholders[idx]
          param_i, param_placeholder_i, param_assign_op_i = weight_placeholders[idx]
          dim_i = param_i.get_shape().as_list()
  
          # recover the masked weights to zeros if they drifted
          param_masked = tf.multiply(mask_i, param_i)
          param_masked_val = sess.run(param_masked)
  
          if config['similarity'] == 'norm_euclidean':
              # first retrieve nonzero rows from the parameter matrix
              row_norm = norm(param_masked_val, axis=1)
              row_norm[row_norm < threshold] = 0 # remove very small row norms
              num_nonzero_rows = np.count_nonzero(row_norm)
              nonzero_row_idx = np.flatnonzero(row_norm)
              nonzero_rows = param_masked_val[row_norm>0, :]
              norm_nonzero_rows = norm(nonzero_rows, axis=1)
              # then compute the normalized Euclidean norm similarity matrix, we should take the negative
              # so that similar rows have a larger affinity value
              similarity_val = np.zeros([num_nonzero_rows, num_nonzero_rows])
              for i in range(num_nonzero_rows):
                  for j in range(num_nonzero_rows):
                      if norm_nonzero_rows[i] > norm_nonzero_rows[j]:
                          similarity_val[i,j] = -norm(nonzero_rows[i,:] - nonzero_rows[j,:])/norm_nonzero_rows[j]
                      else:
                          similarity_val[i,j] = -norm(nonzero_rows[i,:] - nonzero_rows[j,:])/norm_nonzero_rows[i]
  
              #calculate the display similarity matrix without removing the zero valued rows
              num_rows = row_norm.size
              display_similarity_val = np.zeros([num_rows, num_rows])
              display_similarity_val_partial = np.zeros([nonzero_row_idx.size, nonzero_row_idx.size])
  
              zero_row_idx = np.where(row_norm<threshold)[0]
              for k in zero_row_idx:
                  display_similarity_val[k,:] = 0
                  display_similarity_val[:, k] = 0
  
              partial_idx = np.arange(nonzero_row_idx.size)
  
              for idx_newi, i in np.nditer([partial_idx,nonzero_row_idx]):
                  for idx_newj, j in np.nditer([partial_idx,nonzero_row_idx]):
                      if row_norm[i] > row_norm[j]:
                          display_similarity_val[i,j] = display_similarity_val_partial[idx_newi, idx_newj] = np.dot(param_masked_val[i,:], param_masked_val[j,:])/row_norm[i]**2
                      else:
                          display_similarity_val[i,j] = display_similarity_val_partial[idx_newi, idx_newj] = np.dot(param_masked_val[i,:], param_masked_val[j,:])/row_norm[j]**2
  
          # save the similarity val
          if get_group:
              sim_name = config['plot_dir'] + 'epoch_{0}_similarity_{1}.npy'.format(epoch, idx+1)
              np.save(sim_name, display_similarity_val)
  
          #Visualize the similarity matrix
          num_rows = display_similarity_val.shape[0]
          row_idx = np.linspace(1, num_rows, float(num_rows))
  
          x, y = np.meshgrid(row_idx, row_idx)
  
          # plot the grid search results for accuracy
          figName = config['plot_dir'] + 'epoch_{0}_similarity_{1}.png'.format(epoch, idx+1)
          if config['similarity'] == 'norm_euclidean':
              imagesc(x, y, display_similarity_val, 'row_idx', 'row_idx', 'similarity matrix of layer {0}'.format(idx+1),
               False, figName)
  
          similarity_val_arr = np.reshape(similarity_val, [-1])
          display_similarity_val_arr = np.reshape(display_similarity_val, [-1])
  
          # save the nonzero similarity value
          nonzero_display_similarity_val_arr = display_similarity_val_arr[np.abs(display_similarity_val_arr)>0]
          print('Layer {}, the median of the similarity_val is {}'.format(idx, np.median(nonzero_display_similarity_val_arr)))
          if get_group:
              np.savetxt(config['plot_dir'] + 'sim_vec_epoch{}.csv'.format(epoch), display_similarity_val_arr)
          plt.hist(nonzero_display_similarity_val_arr)
          plt.title('Similarity histogram at epoch {}'.format(epoch+1))
          plt.savefig(config['plot_dir'] + 'layer{}_train_sim_{}.jpg'.format(idx, epoch))
          plt.close()
  
          # Cluster the paramter matrix with affinity propagation
          if num_nonzero_rows > 0:
              if config['similarity'] == 'norm_euclidean':
                  af = AffinityPropagation(affinity='precomputed', preference=config['preference']).fit(display_similarity_val_partial)
  
              cluster_centers_indices = af.cluster_centers_indices_
              num_clusters = np.size(cluster_centers_indices)
  
              print("number of clusters {}".format(num_clusters))
  
              # get the labels for the nonzero rows and construct the tuple list for storing the group information. group_info_i is only for the ith layer
              if get_group:
                  labels = af.labels_
                  group_info_i = [tuple()] * num_clusters
                  for i in range(len(labels)):
                      group_info_i[labels[i]] = group_info_i[labels[i]] + (nonzero_row_idx[i], )
  
                  group_info.append(group_info_i)
  
              # put the clustering results in the tuple for return
              num_nonzero_rows_tuple = num_nonzero_rows_tuple + (num_nonzero_rows, )
              num_clusters_tuple = num_clusters_tuple + (num_clusters, )
  
              print('Nonzero rows: {}, Number of Clusters: {}'.format(num_nonzero_rows, num_clusters))
          else:
              num_nonzero_rows_tuple = num_nonzero_rows_tuple + (0, )
              num_clusters_tuple = num_clusters_tuple + (0, )
  
              print('Nonzero rows: {}, Number of Clusters: {}'.format(0, 0))
  
  
  
      return num_nonzero_rows_tuple, num_clusters_tuple, group_info
  
  
  
  def group_averaging(param_val_masked, group_info):
      '''
      This function averages the rows which belongs to the same group
  
      Args:
          param_val_masked: the masked parameter matrix
          group_info: list of tuples, each tuple contains the rows indices that belong to the same group
  
      Return:
          param_val_avg: the averaged parameter rows after enforcing parameter sharing
      '''
      param_val_avg = param_val_masked
  
      for group_i in group_info:
  
          group_i = list(group_i)
          avg_param_i = np.average(param_val_masked[group_i, :], axis=0)
          param_val_avg[group_i, :] = avg_param_i
  
      return param_val_avg
  