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


gnuplotTemplate = file(settings.GNUPLOT_SCATTER_TEMPLATE_FILE, 'r').read()

input_sub_dir = ''
figure_format_list = ['eps']
network_distance_measure_list = ['pca_2']
ignore_methods = []
recalculate_pairwise_distances = False
random_shuffling = 0

opts, args = getopt.getopt(sys.argv[1:], "s:f:d:h", ["subdir=","help","ignore-methods=","random-shuffle=","recalculate-pairwise-distances"])
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 == '-d':
        network_distance_measure_list = a.split(',')
    elif o == '--recalculate-pairwise-distances':
        recalculate_pairwise_distances = True
    elif o == '--ignore-methods':
        ignore_methods = a.split(',')
    elif o  == '--random-shuffle':
        random_shuffling = int(a)
    elif o in ('-h','--help'):
        print """Usage:
%s model_name method_name output_sub_dir
-s,--subdir                         Input/output subdirectory
-f                                  Figure format (e.g., eps or png)
-d                                  Network distance measure (spearman|pearson|pca_x)
--recalculate-pairwise-distances    Recalculate pairwise distances
--ignore-methods                    Methods to ignore (if method name contains any of the input substrings, ignore it)
--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)

model_name_list = args[0].split(',')
output_sub_dir = args[2]


network_name_list = []
#dict of the form method => network instance => graphlet counts
network_graphlet_vector_list = {}
#dict of the form method => pairwise distances
network_distances = {}

for model_name in model_name_list:
    model_dir = settings.BASE_INPUT_DIR + input_sub_dir + 'model/%s/%s/' % (model_utils.getModelType(model_name),model_name)
    #for each model, get all random instances
    network_name_list += [name for name in os.listdir(model_dir) if os.path.isdir(os.path.join(model_dir, name))]


#dictionary needed to support random reshuffling of labels
network_class_dict = {}
for network_name in network_name_list:
    network_class_dict[network_name] = model_utils.getModelClass(network_name)

if args[1] == 'all' or '#' in args[1]: #get all available methods
    i=0
    temp_methods = set()
    for network_name in network_name_list:
        network_methods = set(gdv_utils.getGdvsMethods(network_name,input_sub_dir + 'model/%s/%s/' % (model_utils.getModelType(network_name),model_utils.getModelName(network_name)), only_with_gdvs=False))

        if i>0:
            temp_methods = network_methods.intersection(temp_methods)
        else:
            temp_methods = network_methods


        i+=1
    method_name_list = list(temp_methods)

    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_10')

else:
    method_name_list = args[1].split(',')

if len(ignore_methods)>0:
    print 'Ignoring methods: %s' % ignore_methods
    for m in ignore_methods:
          for method in method_name_list[:]:
              if (m in method) and (method in method_name_list):
                  method_name_list.remove(method)

#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 %s random model networks and %s method(s)' % (len(network_name_list), len(method_name_list))

pca_output_dir = settings.BASE_OUTPUT_DIR + 'MODEL_ANALYSIS/%s/pca/' % output_sub_dir


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


    for distance_measure in network_distance_measure_list:

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

        distances_file = settings.BASE_OUTPUT_DIR + '/MODEL_ANALYSIS/%s/' % (output_sub_dir) + '%s_distances/%s_distances_%s.txt' % (distance_measure,method_name,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:
                distance_list = pca_utils.generate_random_points(len(network_name_list),2) # instead of doing PCA, just generate random points
            else:
                print 'Loading graphlet count vectors...'
                network_graphlet_vector_list = []

                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
                    for i in range(0,len(network_name_list)):
                        network_sub_dir = input_sub_dir + 'model/%s/%s' % (model_utils.getModelType(network_name_list[i]),model_utils.getModelName(network_name_list[i]))
                        temp_vector = gdv_utils.loadGraphletCounts(gdv_utils.getMethodGraphletCountsFile(network_name_list[i],
                                                                                                       method_name.replace('dsgdvs','dgdvs').replace('dcsgdvs','dcgdvs'),network_sub_dir),
                                                                                                                   two_columns=False)
                        temp_vector +=  gdv_utils.loadGraphletCounts(gdv_utils.getMethodGraphletCountsFile(network_name_list[i],
                                                                                                       'snapshots_2_1_stgdvs_3',network_sub_dir),
                                                                                                                   two_columns=True)
                        network_graphlet_vector_list.append(temp_vector)
                else:
                    for i in range(0,len(network_name_list)):
                        network_sub_dir = input_sub_dir + 'model/%s/%s' % (model_utils.getModelType(network_name_list[i]),model_utils.getModelName(network_name_list[i]))
                        network_graphlet_vector_list.append(gdv_utils.loadGraphletCounts(gdv_utils.getMethodGraphletCountsFile(network_name_list[i],method_name,network_sub_dir),
                                                                                                               two_columns=('sgdvs' in method_name or'stgdvs' in method_name)))
                if 'pca' in distance_measure:
                    num_pca_components = int(distance_measure.split('_')[1])

                    pca_values_file = pca_output_dir + '%s_pca_%s.txt' %  (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
                        distance_list = pca_utils.load_pca_file(pca_values_file)
                    else:
                        print 'Calculating PCA...'
                        distance_list = pca_utils.pca_decomposition(network_graphlet_vector_list,network_name_list, pca_values_file, num_pca_components = num_pca_components)

                    pca_utils.draw_pca(distance_list,network_name_list,pca_output_dir + 'plots/', method_name, figure_format_list=figure_format_list, object_to_class_func=model_utils.getModelClassVerbose)
                else:
                    distance_list = network_graphlet_vector_list

            summary_distances_output_dir = settings.BASE_OUTPUT_DIR + '/MODEL_ANALYSIS/%s/%s_distances/' % (output_sub_dir,distance_measure)
            summary_distances_file = summary_distances_output_dir + '%s_%s_summary_distances_%s.txt' % (output_sub_dir,method_name,distance_measure)

            print 'Calculating pairwise distances...'
            distances = pca_utils.calculate_pairwise_distances(distance_list,network_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,
                                                                                    object_to_class_func=model_utils.getModelClass
            )

        print 'Calculating accuracy...'
        pr_roc_output_dir = settings.BASE_OUTPUT_DIR + '/MODEL_ANALYSIS/%s/%s_distances/' % (output_sub_dir,distance_measure) + ('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,
                                                       object_to_class_func = lambda x: network_class_dict[x])

        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):
                network_labels_shuffled = utils.shuffle_dict(network_class_dict)

                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: network_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 distance_measure in network_distance_measure_list:
        print 'Averaging random results for %s...' % distance_measure
        distances_output_dir = settings.BASE_OUTPUT_DIR + '/MODEL_ANALYSIS/%s/%s_distances/' % (output_sub_dir,distance_measure)
        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)









