import sys, getopt, os, subprocess, math
import numpy as np
import scipy.stats, scipy.spatial
import utils, gdv_utils, plot_utils, pr_roc_utils, draw_utils, model_utils, ground_truth_utils, pca_utils
import settings

input_sub_dir = ''
figure_format_list = ['eps']
distance_measure_list = ['pca_3']
ignore_pairs_ground_truth_value = None
min_node_degree = None
recalculate_pairwise_distances = False
random_shuffling = 0

opts, args = getopt.getopt(sys.argv[1:], "s:f:d:h", ["subdir=","help","recalculate-pairwise-distances","ignore-pairs-ground-truth=","min-node-degree=","random-shuffle="])
for o, a in opts:
    if o in ('-s', '--subdir'):
        input_sub_dir = a  + '/'
    elif o in ('-f'):
        figure_format_list = a.split(',')
    elif o == '--recalculate-pairwise-distances':
        recalculate_pairwise_distances = True
    elif o in ('--ignore-pairs-ground-truth'):
        ignore_pairs_ground_truth_value = a
    elif o == '-d':
        distance_measure_list = a.split(',')
    elif o == '--min-node-degree':
         min_node_degree = int(a)
    elif o  == '--random-shuffle':
         random_shuffling = int(a)
    elif o in ('-h','--help'):
        print """Usage:
%s network_name method_name ground_truth_type
-s,--subdir                         Input/output subdirectory
--recalculate-pairwise-distances    Recalculate pairwise distances
--ignore-pairs-ground-truth         If two objects have this class - exclude this pair from analysis
-d                                  Network distance measure (spearman|pearson|pca)
--min-node-degree                   Remove nodes with small degrees from consideration
--random-shuffle                    For each method perform random shuffling of labels to obtain another null model
-h,--help                           Print this""" % sys.argv[0]
        sys.exit(0)

network_name = args[0]
ground_truth_type_list = args[2].split(',')

pca_output_dir = settings.BASE_OUTPUT_DIR + '%s%s/pca/' % (input_sub_dir,network_name)

if args[1] == 'all' or '#' in args[1]: #get all available methods
    method_name_list = list(set(gdv_utils.getGdvsMethods(network_name,input_sub_dir, only_with_gdvs=False)))

    if not args[1] == 'all':
        if args[1].startswith('#') and not args[1].endswith('#'):
            method_name_list = [m for m in method_name_list if m.endswith(args[1][1:])]
        elif not args[1].startswith('#') and args[1].endswith('#'):
            method_name_list = [m for m in method_name_list if m.startswith(args[1][:-1])]
        elif args[1].startswith('#') and args[1].endswith('#'):
            method_name_list = [m for m in method_name_list if args[1][1:-1] in m]
        else:
            raise Exception('Unknown method identifier: %s' % args[1])
    else:
        method_name_list.append('random_5')

elif args[1].find('key_')==0:
    method_name_list = gdv_utils.getMethodsPredefinedSets(args[1])
else:
    method_name_list = args[1].split(',')


#handle 'random' methods: (random_5 -> random_1,random_2,random_3,random_4,random_5)
for i in range(0,len(method_name_list)):
    if method_name_list[i].find('random_') == 0:
        values = method_name_list[i].split('_')
        random_methods = []
        for j in range(1,int(values[1])):#the last random method is already there
            random_methods.append('random_%s' % j)
        method_name_list += random_methods
        break

method_name_list = sorted(method_name_list, cmp = gdv_utils.compareMethodNames)
print 'Analyzing the following %s methods: %s' % (len(method_name_list),', '.join(method_name_list))

i_method = 1
for method_name in method_name_list:
    print '%s/%s. %s' % (i_method,len(method_name_list), method_name)
    i_method += 1
    is_random_clustering = method_name.find('random')==0


    i_ground_truth = 1
    for ground_truth_type in ground_truth_type_list:

        print '%s/%s. %s' % (i_ground_truth,len(ground_truth_type_list), ground_truth_type)
        i_ground_truth += 1
        node_labels = ground_truth_utils.get_node_labels(network_name,ground_truth_type)


        for distance_measure in distance_measure_list:

            print 'Using distance measure %s' % (distance_measure)

            distances_output_dir = pca_output_dir
            distances_file = distances_output_dir + ('' if not ignore_pairs_ground_truth_value else (ground_truth_type + '/'))\
                             + '%s_distances%s_%s.txt' % (method_name,'' if not ignore_pairs_ground_truth_value else '_ignoring',distance_measure)


            if not is_random_clustering and not recalculate_pairwise_distances and os.path.exists(distances_file):
                print 'Loading pairwise distances from %s...' % distances_file
                distances = pca_utils.load_pairwise_distances(distances_file)
            else:
                print 'Calculating pairwise distances...'

                if is_random_clustering:
                    node_name_list = gdv_utils.getNodeNames(network_name, input_sub_dir, min_node_degree)
                    distance_list = pca_utils.generate_random_points(len(node_name_list),3) # instead of doing PCA, just generate random points
                else:
                    if 'pca' in distance_measure:
                        num_pca_components = int(distance_measure.split('_')[1])

                        print 'Loading GDVs...'

                        if 'dsgdvs' in method_name or 'dcsgdvs' in method_name: #combined dynamic and static-temporal graphlet counts - to capture information both between and inside snapshots
                            d_gdvFile = gdv_utils.getMethodGdvFile(network_name,method_name.replace('dsgdvs','dgdvs').replace('dcsgdvs','dcgdvs'),input_sub_dir)
                            d_node_name_list, d_gdv_list = gdv_utils.loadGdvFile(d_gdvFile)
                            st_gdvFile = gdv_utils.getMethodGdvFile(network_name,'stgdvs_4',input_sub_dir)
                            st_node_name_list, st_gdv_list = gdv_utils.loadGdvFile(st_gdvFile)

                            print 'Vector lengths: %s and %s' % (len(d_gdv_list[0]),len(st_gdv_list[0]))
                            #node name lists may be in different order
                            gdv_list = []
                            for i in xrange(len(d_node_name_list)):
                                gdv_list.append(st_gdv_list[st_node_name_list.index(d_node_name_list[i])][:] + d_gdv_list[i][:])
                            node_name_list = d_node_name_list[:]

                            #print st_node_name_list
                            #print node_name_list

                           # for i in xrange(len(st_gdv_list)):
                            #    if st_gdv_list[i] != gdv_list[i]:
                           #         print 'AAAA'
                            #gdvFile = gdv_utils.getMethodGdvFile(network_name,'snapshots_2month_1_stgdvs_4',input_sub_dir)
                            #node_name_list, gdv_list = gdv_utils.loadGdvFile(gdvFile)
                        else:
                            gdvFile = gdv_utils.getMethodGdvFile(network_name,method_name,input_sub_dir)
                            node_name_list, gdv_list = gdv_utils.loadGdvFile(gdvFile)


                        if min_node_degree is not None:
                            temp_node_names = gdv_utils.getNodeNames(network_name, input_sub_dir, min_node_degree)
                            delete_nodes = list(set(node_name_list) - set(temp_node_names))
                            print 'Removing all nodes with degree less than %s from consideration (%s such nodes)...' % (min_node_degree, len(delete_nodes))
                            for v in delete_nodes:
                                i = node_name_list.index(v)
                                del node_name_list[i]
                                del gdv_list[i]

                        pca_values_file = pca_output_dir + '%s_%s_pca_%s.txt' %  (network_name, method_name, num_pca_components)

                        if not recalculate_pairwise_distances and os.path.exists(pca_values_file):
                            print 'Loading PCA values from %s...' % pca_values_file
                            pca_list = pca_utils.load_pca_file(pca_values_file)
                        else:
                            print 'Calculating PCA...'
                            pca_list = pca_utils.pca_decomposition(gdv_list, node_name_list, pca_values_file, num_pca_components = num_pca_components)

                        #to fix a bug for aging network and stgdvs
                        if len(node_labels.keys()) < len(node_name_list):
                            print 'WARNING: Too many nodes in gdv file (more than in the aggregated_1 network): %s > %s - deleting extra nodes...' % (len(node_name_list),len(node_labels.keys()))
                            temp_pca_list = list(pca_list)
                            i=0
                            while True:
                                if not node_labels.has_key(node_name_list[i]):
                                    print 'Deleting %s - %s' % (i, node_name_list[i])
                                    del node_name_list[i]
                                    del gdv_list[i]
                                    del temp_pca_list[i]
                                    i-=1
                                i+=1
                                if i >= len(node_name_list):
                                    break
                            pca_list = np.array(temp_pca_list)
                            del temp_pca_list
                            print 'New number of nodes: %s (%s nodes in ground truth)' % (len(node_name_list),len(node_labels.keys()))


                        if not is_random_clustering:
                            #TODO: add number of PCA components to the filename
                            pca_utils.draw_pca(pca_list,
                                               node_name_list,
                                               pca_output_dir +  '%s/plots/' % ground_truth_type, method_name, figure_format_list=figure_format_list,
                                               object_to_class_func = lambda x: ground_truth_utils.get_node_label_title(network_name,ground_truth_type,node_labels[x]))

                        distance_list = pca_list

                summary_distances_output_dir = pca_output_dir + '%s/%s_distances%s/' % (ground_truth_type,distance_measure,'' if not ignore_pairs_ground_truth_value else '_ignoring')
                summary_distances_file = summary_distances_output_dir + '%s_%s_summary_distances_%s.txt' % (network_name,method_name,distance_measure)

                print 'Calculating pairwise distances...'
                distances = pca_utils.calculate_pairwise_distances(distance_list,
                                                                   node_name_list,
                                                                   distance_measure = distance_measure,
                                                                   sort_distances=True,
                                                                   output_file = distances_file if not is_random_clustering else None,
                                                                   summary_distances_file = summary_distances_file if not is_random_clustering else None,
                                                                   ignore_pairs_ground_truth_value=ignore_pairs_ground_truth_value,
                                                                   object_to_class_func= lambda x: ground_truth_utils.get_node_label_title(network_name,ground_truth_type,node_labels[x])
                )

            print 'Calculating accuracy...'
            pr_roc_output_dir = pca_output_dir + '%s/%s_distances%s/' % (ground_truth_type,distance_measure,'' if not ignore_pairs_ground_truth_value else '_ignoring') + ('random/' if is_random_clustering else '')

            pr_roc_utils.calculate_class_separation_pr_roc(distances, pr_roc_output_dir + '%s_pr_roc.txt' % method_name,
                                                           ignore_pairs_ground_truth_value=ignore_pairs_ground_truth_value,
                                                           object_to_class_func= lambda x: ground_truth_utils.get_node_label_title(network_name,ground_truth_type,node_labels[x]),
                                                           #increment_indices = range(100,10100,100),#by 100 pairs
                                                           verbose = False)

            if random_shuffling !=0 and not is_random_clustering:
                print 'Randomly shuffling labels %s times...' % random_shuffling
                file_list = []
                for i in range(0,random_shuffling):
                    node_labels_shuffled = utils.shuffle_dict(node_labels)

                    output_file = pr_roc_output_dir + 'random_shuffle/%s_%s_pr_roc.txt' % (method_name,i+1)
                    pr_roc_utils.calculate_class_separation_pr_roc(distances,output_file,
                                                               object_to_class_func= lambda x: ground_truth_utils.get_node_label_title(network_name,ground_truth_type,node_labels_shuffled[x]),
                                                               verbose=False)
                    file_list.append(output_file)

                #average_results
                print 'Averaging results...'
                pr_roc_utils.average_pr_roc(file_list,pr_roc_output_dir + '%s_random_shuffle_pr_roc.txt' % method_name)

#Average random results
any_random_methods = len([m for m in method_name_list if m.find('random')==0])>0
if any_random_methods:
    for ground_truth_type in ground_truth_type_list:
        for distance_measure in distance_measure_list:
            print 'Averaging random results for %s and %s distance...' % (ground_truth_type, distance_measure)
            distances_output_dir = pca_output_dir + '%s/%s_distances%s/' % (ground_truth_type, distance_measure, '' if not ignore_pairs_ground_truth_value else '_ignoring')
            file_list = []
            files_dir = distances_output_dir + 'random/'
            for f in os.listdir(files_dir):
                if 'random' in f and 'pr_roc' in f:
                    file_list.append(files_dir + f)

            aggregate_filename = distances_output_dir +'random_pr_roc.txt'
            pr_roc_utils.average_pr_roc(file_list,aggregate_filename)