import sys, getopt, os, subprocess, math
import numpy as np
import scipy.stats
import utils, gdv_utils, plot_utils, pr_roc_utils, ground_truth_utils
import clustering
import networkx as nx
import settings


inputSubDir = ''
n_clusters_list = None
kmedoids_npass = 100
p_value_threshold = 0.05
run_random_clustering=0

opts, args = getopt.getopt(sys.argv[1:], "s:k:p:g:r:h", ["subdir=","help","kmedoids-npass="])
for o, a in opts:
    if o in ('-s', '--subdir'):
        inputSubDir = a  + '/'
    elif o in ('-k'):
        n_clusters_list = sorted([int(x) for x in a.split(',')])
    elif o in ('-p'):
        p_value_threshold = float(a)
    elif o in ('--kmedoids-npass'):
        kmedoids_npass = int(a)
    elif o in ('-r'):
        run_random_clustering = int(a)
    elif o in ('-h','--help'):
        print """Usage:
%s network_name method_name ground_truth_type
-s,--subdir      Input/output subdirectory
-k               Number of clusters
-p               p-value threshold
--kmedoids-npass Number of repetitions of k-medoids
-r               Generate also results for random clusterings with the same cluster sizes (number of random runs)
-h,--help        Print this""" % sys.argv[0]
        sys.exit(0)

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


if n_clusters_list is None:
    n_clusters_list = [1000] if network_name == 'aging' else [50]

m_counter = 1
for method_name in method_name_list:
    print '%s/%s. %s_%s' % (m_counter,len(method_name_list), network_name, method_name)
    m_counter += 1

    dist_matrix_file = os.path.dirname(gdv_utils.getMethodGdvFile(network_name,method_name)) + '/%s_%s_dist_matrix.txt' % (network_name,method_name)
    if not os.path.exists(dist_matrix_file):
        print 'Creating and saving filtered distance matrix'
        #using a separate script since otherwise running kmedoids leads to segmentation fault
        p = subprocess.Popen(['python', 'create_gdvs_distance_matrix.py',network_name,method_name], stdout=subprocess.PIPE)
        p_output = p.stdout.read()
        if p_output != '':
            print(p_output)


    print 'Loading filtered distance matrix'
    dist_matrix = np.loadtxt(dist_matrix_file, skiprows=1)
    print 'Loaded distance matrix %sx%s' % (dist_matrix.shape[0],dist_matrix.shape[1])
    node_names = open(dist_matrix_file).readline().split()[1:]

    N = len(node_names) #total number of nodes

    if N != dist_matrix.shape[0] or N != dist_matrix.shape[1]:
        raise Exception('%s != %s and %s' % (N, dist_matrix.shape[0], dist_matrix.shape[1]))

    k_counter = 1
    for n_clusters in n_clusters_list:
        print '%s/%s. %s_%s - %s/%s. %s clusters' % (m_counter-1,len(method_name_list), network_name, method_name, k_counter,len(n_clusters_list), n_clusters)
        k_counter += 1

        clustering_dir = settings.BASE_OUTPUT_DIR + inputSubDir + '%s/clusterings/' % (network_name)
        clustering_file= clustering_dir + '%s_%s_km_%s_clustering.txt' % (network_name, method_name, n_clusters)

        if not os.path.exists(clustering_file):

            print 'Performing clustering...'

            clustering_list, error, nfound = clustering.kmedoids(dist_matrix, n_clusters, npass = kmedoids_npass)
            #clustering_list = clustering.generate_random_clustering(N, 1000), error = 0, nfound = 0
            print 'Finished running KM'
            clustering_dict_temp = clustering.convert_clustering_to_dict_format(clustering_list)
        
            #convert elements in clustering dict from v to node_names[v]
            clustering_dict = dict()
            for c_key in clustering_dict_temp:
                clustering_dict[c_key] = [node_names[v] for v in clustering_dict_temp[c_key]]


            if not os.path.exists(clustering_dir):
                os.makedirs(clustering_dir)

            #save to file
            clustering.saveToFile(clustering_dict, clustering_file)

            #draw histogram of cluster sizes
            plot_utils.draw_values_histogram(clustering.get_cluster_sizes(clustering_dict), clustering_dir+'cluster_sizes_plots/','%s_%s_cluster_sizes' % (network_name,method_name),
                                             {'prop_group_by':2, 'x_label': 'Cluster sizes'})
        else:
            print 'Using existing clustering'
            clustering_dict = clustering.loadFromFile(clustering_file)
            error = 0
            nfound = 0


        ## make a list of the original clustering and the random clustering
        org_rd_clustering_list=dict()
        org_rd_clustering_list.update({network_name:clustering_dict}) 
        error_list={network_name:error}
        nfound_list={network_name:nfound}


        if run_random_clustering > 0:
            print 'Generating results for random clustering with the same cluster sizes'
            if not os.path.exists(clustering_dir + 'random/'):
                os.makedirs(clustering_dir + 'random/')

            for i_random in range(1,run_random_clustering+1):
                # check whether the random clustering file already exists
                rd_clustering_file= clustering_dir + 'random/%s_random_%s_%s_km_%s_clustering.txt' % (network_name,i_random, method_name, n_clusters)



                #convert elements in clustering dict from v to node_names[v]
                random_clustering_dict = dict()
                #Always generate new random clustering
                #if not os.path.isfile(rd_clustering_file):
                print 'Performing random clustering #%s...' % i_random
                print 'Random clustering: %s_random_%s_%s_km_%s_clustering.txt' % (network_name, i_random, method_name, n_clusters)
                cluster_sizes=clustering.get_cluster_sizes(clustering_dict)
                random_clustering_list=clustering.generate_random_clustering(N,n_clusters,cluster_sizes)
                ## convert to dict format
                rd_clustering_dict_temp = clustering.convert_clustering_to_dict_format(random_clustering_list)


                for c_key in rd_clustering_dict_temp:
                    random_clustering_dict[c_key] = [node_names[v] for v in rd_clustering_dict_temp[c_key]]

                #save to file
                clustering.saveToFile(random_clustering_dict, rd_clustering_file)

                #draw histogram of cluster sizes
                plot_utils.draw_values_histogram(clustering.get_cluster_sizes(random_clustering_dict), clustering_dir+'cluster_sizes_plots/','random_%s_%s_%s_cluster_sizes' % (i_random,network_name,method_name),
                                                 {'prop_group_by':2, 'x_label': 'Cluster sizes'})
                #else:
                #    print 'Random clustering already exists (%s)' % rd_clustering_file
                #    random_clustering_dict = clustering.loadFromFile(rd_clustering_file)

                random_error=0
                random_nfound=0
                ## N= total number of nodes; n_cluster=number of clusters; cluster_sizes=distribution of cluster sizes
                rd_net_name=network_name + '_random_%s' % i_random
                org_rd_clustering_list.update({rd_net_name:random_clustering_dict})
                error_list.update({rd_net_name:random_error})
                nfound_list.update({rd_net_name:random_nfound})


        g_counter = 1
        for ground_truth_type in ground_truth_type_list:
            print '%s/%s. %s_%s - %s/%s. %s clusters - %s/%s. %s' % (m_counter-1,len(method_name_list), network_name, method_name, k_counter-1,len(n_clusters_list), n_clusters,
                                                                                 g_counter,len(ground_truth_type_list), ground_truth_type)
            g_counter += 1

            if network_name == 'aging':
                node_labels = ground_truth_utils.get_aging_node_labels(ground_truth_type)
            elif network_name == 'enron':
                node_labels = ground_truth_utils.get_enron_node_labels(ground_truth_type)

            A = len([v for v in node_labels.keys() if node_labels[v]==1 and v in node_names]) #number of active nodes
            print '%s nodes total, %s nodes are \'active\'' % (N,A)

            output_dir = settings.BASE_OUTPUT_DIR + inputSubDir + '%s/enrichment/%s_%s/' % (network_name, ground_truth_type, n_clusters)

            ## run through all clusterings
            i_clustering_counter = 1

            for network_name_key in org_rd_clustering_list.keys():
                clustering_dict_tmp=org_rd_clustering_list[network_name_key]


                ierror=error_list[network_name_key]
                infound=nfound_list[network_name_key]

                network_name_key_description = 'Computed clustering' if network_name_key == network_name else 'Random clustering (%s)' % network_name_key
                is_random_clustering = network_name_key != network_name


                print '%s/%s. %s_%s - %s/%s. %s clusters - %s/%s. %s  - %s/%s. %s' % (m_counter-1,len(method_name_list), network_name, method_name, k_counter-1,len(n_clusters_list), n_clusters,
                                                                                 g_counter-1,len(ground_truth_type_list), ground_truth_type,
                                                                                 i_clustering_counter,len(org_rd_clustering_list), network_name_key_description)
                i_clustering_counter += 1

                print 'Clustering: %s clusters (%s - max size, %s  - min size, %.3f - average, %.3f - std dev), error=%s, nfound=%s' \
                % (len(clustering_dict_tmp.keys()),len(clustering.get_largest_cluster(clustering_dict_tmp)),len(clustering.get_smallest_cluster(clustering_dict_tmp)),
                 np.average(clustering.get_cluster_sizes(clustering_dict_tmp)), np.std(clustering.get_cluster_sizes(clustering_dict_tmp)), ierror, infound)



                if not os.path.exists(output_dir + 'log/'):
                    os.makedirs(output_dir + 'log/')

                if not os.path.exists(output_dir + 'random/log/'):
                    os.makedirs(output_dir + 'random/log/')

                ##check whether log file exists
                log_file_tmp=output_dir +  ('random/' if is_random_clustering else '')+'log/%s_%s_log.txt' % (network_name_key,method_name)
                if os.path.isfile(log_file_tmp) and False:
                    print 'log file \'%s\' already exists' % log_file_tmp
                else:
                    log_file = open(log_file_tmp , 'wb')

                    log_file.write('%s nodes total, %s nodes are \'active\'\n' % (N,A))
                    log_file.write('Clustered data: %s clusters (%s - max size, %s  - min size, %.3f - average, %.3f - std dev), error=%s, nfound=%s\n' \
                      % (len(clustering_dict_tmp.keys()),len(clustering.get_largest_cluster(clustering_dict_tmp)),len(clustering.get_smallest_cluster(clustering_dict_tmp)),
                         np.average(clustering.get_cluster_sizes(clustering_dict_tmp)), np.std(clustering.get_cluster_sizes(clustering_dict_tmp)), ierror, infound))
                    log_file.write('Cluster sizes: %s\n' % clustering.get_cluster_sizes(clustering_dict_tmp))
                    log_file.write('Clustering: %s\n' % clustering.get_clustering_description_string(clustering_dict_tmp))

                    log_file.write('Cluster enrichments of clusters of size at least 3:\n')


                #draw clustering
                #print 'Drawing clustering...'
                #plot_utils.draw_clustering(clustering_dict,output_dir+'log/clustering_plots/','%s_%s_clustering' % (network_name,method_name), ground_truth_clustering=node_labels, ground_truth_clustering_colors = {0: 'red', 1: 'green'})

                K_total = 0
                i_cluster = 1
                #number of active nodes in clusters of size
                for c_key in clustering_dict_tmp:
                    c = clustering_dict_tmp[c_key]
                    C = len(c) #number of nodes in a cluster
                    if C<3:

                        continue
                    A_cluster = len([v for v in c if node_labels.has_key(v) and node_labels[v]==1])
                    K_total += A_cluster # number of active nodes in a cluster
                    log_file.write('\tCluster%s: active nodes - %s/%s (%s%%)\n' % (i_cluster, A_cluster, C,100.0 * A_cluster/C))
                    i_cluster += 1

                log_file.write('Active nodes in clusters of size at least 3: %s/%s (%.0f%%)\n' % (K_total,A, 100.0 * K_total/A))
                print 'Active nodes in clusters of size at least 3: %s/%s (%.0f%%)' % (K_total,A, 100.0 * K_total/A)


                max_fscore = 0
                k_max_fscore = 0

                #for area under the curve
                precisions = []
                recalls = []
                tp_rates = []
                fp_rates = []


                print 'Calculating accuracy...'
                with open(output_dir +  ('random/' if is_random_clustering else '') + '%s_%s_pr_roc.txt' % (network_name_key,method_name) , 'wb') as outfile:
                    outfile.write("k\tTP\tFP\tTN\tFN\tRecall\tPrecision\tTP rate\tFP rate\tF-score\tAccuracy\n")

                    for k in range(0, 101)[::-1]:
                        tp = 0
                        fp = 0
                        fn = 0
                        tn = 0
                        for c_key in clustering_dict_tmp:
                            c = clustering_dict_tmp[c_key]
                            C = len(c) #number of nodes in a cluster
                            if C<3:
                                continue

                            K = len([v for v in c if node_labels.has_key(v) and node_labels[v]==1]) # number of active nodes in a cluster

                            for v in c:
                                if not node_labels.has_key(v):
                                    continue
                                v_label = node_labels[v]
                                C_hidden = C - 1
                                K_hidden = K-1 if v_label == 1 else K
                                Q_hidden = 100.0 * float(K_hidden)/C_hidden
                                p_value_hidden = 1 - scipy.stats.hypergeom.cdf(K_hidden-1,N,A,C_hidden)

                                if Q_hidden>=k and p_value_hidden <= p_value_threshold:
                                    if v_label == 1:
                                        tp += 1
                                    elif v_label == 0:
                                        fp += 1
                                elif v_label == 1:
                                    fn += 1

                        if tp + fn == 0:
                            recall = 0
                        else:
                            recall = float(tp)/(tp+fn)
                        if tp + fp == 0:
                            precision = 1.0
                        else:
                            precision = float(tp)/(tp+fp)

                        tp_rate = recall
                        if fp + tn == 0:
                            fp_rate = 0
                        else:
                            fp_rate = float(fp)/(fp+tn)

                        if precision + recall == 0:
                            f_score = 0
                        else:
                            f_score = (2.0 * precision * recall) / (precision + recall)

                        if tp + fp + tn + fn == 0:
                            accuracy = 0
                        else:
                            accuracy = float(tp + tn) / (tp + fp + tn + fn)


                        if f_score > max_fscore:
                            max_fscore = f_score
                            k_max_fscore = k

                        outfile.write(str(k/100.0) + "\t" + str(tp) + "\t" + str(fp) + "\t" + str(tn) + "\t" + str(fn) + "\t" + str(recall) + "\t" + str(precision)
                                      + "\t" + str(tp_rate) + "\t" + str(fp_rate)
                                      + "\t" + str(f_score) + "\t" + str(accuracy) + "\n")

                        #print 'k=%03d: tp=%d, fp=%d, tn=%d, fn=%d, precision=%.3f, recall=%.3f, fp-rate=%.3f, f-score=%.3f, accuracy=%.3f' % (k,tp,fp,tn,fn,precision, recall, fp_rate, f_score,accuracy)
                        sys.stdout.write('\rk=%s...' % k)
                        sys.stdout.flush()

                        precisions.append(precision)
                        recalls.append(recall)
                        tp_rates.append(tp_rate)
                        fp_rates.append(fp_rate)


                print '\nMax F-score = %s for k = %s' % (max_fscore, k_max_fscore)
                print 'AUPR  = %.3f' % pr_roc_utils.area_under_the_curve(recalls,precisions)
                print 'AUROC = %.3f' % pr_roc_utils.area_under_the_curve(fp_rates,tp_rates)

                log_file.write('AUPR  = %.3f\n' % pr_roc_utils.area_under_the_curve(recalls,precisions))
                log_file.write('AUROC = %.3f\n' % pr_roc_utils.area_under_the_curve(fp_rates,tp_rates))
                log_file.close()

            if len(org_rd_clustering_list.keys()) > 1:
                print 'Averaging random results...'
                file_list = []
                files_dir = settings.BASE_OUTPUT_DIR + inputSubDir + '%s/enrichment/%s_%s/random/' % (network_name,ground_truth_type,n_clusters)
                for f in os.listdir(files_dir):
                    if 'random' in f and method_name in f:
                        file_list.append(files_dir + f)

                aggregate_filename = settings.BASE_OUTPUT_DIR + inputSubDir + '%s/enrichment/%s_%s/%s_random_%s_pr_roc.txt' % (network_name,ground_truth_type, n_clusters, network_name,method_name)
                pr_roc_utils.average_pr_roc(file_list,aggregate_filename)


            



		

