import os, getopt, sys
import numpy as np
import networkx as nx
import gdv_utils
import settings

def filter_distance_matrix(dist_matrix, node_names, remove_node_names):
    del_indices = [node_names.index(r_node) for r_node in remove_node_names]


    # dist_matrix1 = np.delete(dist_matrix, del_indices, 0)
    # dist_matrix1 = np.delete(dist_matrix1, del_indices, 1)

    keep_indices = sorted(set(range(len(node_names))) - set(del_indices))
    dist_matrix = dist_matrix[keep_indices,:][:,keep_indices]

    # if not np.array_equal(dist_matrix,dist_matrix1):
    #     print '%s and %s' % (dist_matrix.shape, dist_matrix1.shape)
    #     print dist_matrix[:5,:][:,:5]
    #     print dist_matrix1[:5,:][:,:5]
    #     print 'AA'
    #     print dist_matrix[1000:1005,:][:,1000:1005]
    #     print dist_matrix1[1000:1005,:][:,1000:1005]
    #     raise Exception('AAA')


    node_names = np.delete(node_names,del_indices).tolist()
    return dist_matrix, node_names




opts, args = getopt.getopt(sys.argv[1:], "h", ["help"])
for o, a in opts:
    if o in ('-h','--help'):
        print """
Generates distance matrix based gdvs (and removes vertices with low degree)
Usage:
%s network_name method_name
-h,--help        Print this""" % sys.argv[0]
        sys.exit(0)

network_name = args[0]
method_name = args[1]

dist_matrix_file = os.path.dirname(gdv_utils.getMethodGdvFile(network_name,method_name)) + '/%s_%s_dist_matrix.txt' % (network_name,method_name)

print 'Loading node similarities matrix for %s...' % method_name
dist_matrix, node_names = gdv_utils.getMethodGdvsMatrix(network_name,method_name,as_a_distance_matrix=True)
print 'Loaded distance matrix %sx%s' % (dist_matrix.shape[0],dist_matrix.shape[1])

print 'num of nodes: '+str(len(node_names))
print ', '.join(node_names[:10])

remove_node_names = gdv_utils.getNodeNames(network_name,'',4)
print 'Removing %s nodes with degree <4...' % len(remove_node_names)

all_node_names = gdv_utils.getNodeNames(network_name)
extra_nodes  = list(set(node_names) - set(all_node_names))

if len(extra_nodes)>0:
    print '%s_%s contains %s extra nodes comparing to %s_aggregated_1 - removing them' % (network_name,method_name,len(extra_nodes),network_name)
    remove_node_names += extra_nodes

dist_matrix, node_names = filter_distance_matrix(dist_matrix,node_names,remove_node_names)

print 'Saving distance matrix...'
#can just use np.savetxt(dist_matrix_file, dist_matrix, header='\t'.join(node_names)) in newer versions
headersTmp='\t'.join(node_names)
headers='#\t'+headersTmp+'\n'
np.savetxt(dist_matrix_file+"tmp", dist_matrix)
with open(dist_matrix_file+"tmp") as f1:
    with open(dist_matrix_file, "w") as f2:
        f2.write(headers)
        for line in f1:
          f2.write(line)
    os.remove(dist_matrix_file+"tmp")
