import os, math, sys
import settings
import numpy as np

def area_under_the_curve(x_vals, y_vals):
    area_sum = 0
    for i in range(0, len(x_vals)-1):
        area_sum += (x_vals[i+1] - x_vals[i])*(y_vals[i]+y_vals[i+1])*.5
    return area_sum


def get_pr_roc_file_column_number(label, starting_from_one = True, contains_std_dev = False):
    c = 0
    if label == 'p':
        c =  7
    elif label == 'r':
        c =  6
    elif label == 'tp':
        c =  2
    elif label == 'fp':
        c =  3
    elif label == 'tp_rate':
        c =  8
    elif label == 'fp_rate':
        c =  9
    elif label == 'f':
        c =  10
    elif label == 'k':
        c =  1
    else:
        c =  0

    if contains_std_dev:
        c = 2*(c-1) if c > 1 else 1

    if starting_from_one:
        return c
    else:
        return c-1

def get_label_title(label):
    if label == 'p':
        return 'Precision'
    elif label == 'r':
        return 'Recall'
    elif label == 'tp_rate':
        return 'TP-rate'
    elif label == 'fp_rate':
        return 'FP-rate'
    elif label == 'f':
        return 'F-score'
    elif label == 'k':
        return 'k'
    return label

def get_AUPR(pr_roc_file, contains_std_dev = False):
    """
    Calculate AUPR from Precision/Recall file
    :param pr_roc_file:
    :param contains_std_dev:
    :return:
    """
    if os.path.isfile(pr_roc_file) == False:
      print pr_roc_file+' does not exists'
      return None
    lines = open(pr_roc_file).readlines()
    precisions = []
    recalls = []

    p_col = get_pr_roc_file_column_number('p', False, contains_std_dev = contains_std_dev)
    r_col = get_pr_roc_file_column_number('r', False, contains_std_dev = contains_std_dev)

    for l in lines[1:]:
        values = l.split('\t')
        precisions.append(float(values[p_col]))
        recalls.append(float(values[r_col]))
    return area_under_the_curve(recalls,precisions)

def get_AUROC(pr_roc_file, contains_std_dev = False):
    """
    Calculate AUROC from Precision/Recall file
    :param pr_roc_file:
    :param contains_std_dev:
    :return:
    """
    if os.path.isfile(pr_roc_file) == False:
      print pr_roc_file+' does not exists'
      return None
    lines = open(pr_roc_file).readlines()
    fp_rates = []
    tp_rates = []

    fp_rate_col = get_pr_roc_file_column_number('fp_rate', False, contains_std_dev = contains_std_dev)
    tr_rate_col = get_pr_roc_file_column_number('tp_rate', False, contains_std_dev = contains_std_dev)

    for l in lines[1:]:
        values = l.split('\t')
        fp_rates.append(float(values[fp_rate_col]))
        tp_rates.append(float(values[tr_rate_col]))
    return area_under_the_curve(fp_rates,tp_rates)

def get_max_fscore(pr_roc_file, contains_std_dev = False):
    if os.path.isfile(pr_roc_file) == False:
       print pr_roc_file+' does not exists'
       return None
    lines = open(pr_roc_file).readlines()

    fscore_col = get_pr_roc_file_column_number('f', False, contains_std_dev = contains_std_dev)
    k_col = get_pr_roc_file_column_number('k', False, contains_std_dev = contains_std_dev)
    max_fscore = -1
    max_fscore_k = -1

    for l in lines[1:]:
        values = l.split('\t')
        f_score = float(values[fscore_col])
        k = float(values[k_col])

        if f_score >= max_fscore:
            max_fscore = f_score
            max_fscore_k = k

    return max_fscore,max_fscore_k



def get_break_even_point(pr_roc_file, contains_std_dev = False):
    """
    Calculate break-even point
    :param pr_roc_file:
    :param contains_std_dev:
    :return: Break-even point, corresponding value of k
    """
    if os.path.isfile(pr_roc_file) == False:
      print pr_roc_file+' does not exists'
      return None
    lines = open(pr_roc_file).readlines()

    precision_col = get_pr_roc_file_column_number('p', False, contains_std_dev = contains_std_dev)
    recall_col = get_pr_roc_file_column_number('r', False, contains_std_dev = contains_std_dev)
    k_col = get_pr_roc_file_column_number('k', False, contains_std_dev = contains_std_dev)

    p_prev = None
    r_prev = None
    k_prev = None

    for l in lines[1:]:
        values = l.split('\t')
        p = float(values[precision_col])
        r = float(values[recall_col])
        k = float(values[k_col])
        if p == r:
            return p,k

        if p_prev is not None and (p_prev - r_prev) * (p - r) < 0:
            p = seg_intersect([k_prev,p_prev],[k,p],[k_prev,r_prev],[k,r])
            return p[1],p[0]
        p_prev = p
        r_prev = r
        k_prev = k

    return None

def seg_intersect(a1,a2, b1,b2):
    """
    Intersection of two line segments
    :param a1:
    :param a2:
    :param b1:
    :param b2:
    :return:
    """
    a1 = np.array(a1)
    a2 = np.array(a2)
    b1 = np.array(b1)
    b2 = np.array(b2)

    da = a2-a1
    db = b2-b1
    dp = a1-b1
    dap = np.empty_like(da)
    dap[0] = -da[1]
    dap[1] = da[0]

    denom = np.dot(dap, db)
    num = np.dot(dap, dp)
    return (num / denom) * db + b1



def average_pr_roc(file_list,output_filename):
    aggregate_file = open(output_filename, 'wb')
    file_lines = open(file_list[0]).readlines()
    headings = file_lines[0].replace('\n','').split('\t')


    values = []
    for k in range(0, len(file_lines)-1):
        values.append([])
        for heading in headings:
            values[k].append([])

    for input_filename in file_list:
        file = open(input_filename, 'r')
        lines = file.readlines()
        lines.pop(0)
        for k in range(0, len(lines)):
            terms = lines[k].split('\t')
            for col in range(0, len(terms)):
                values[k][col].append(float(terms[col]))
        file.close()

    aggregate_file.write('k\t')
    for heading in headings[1:]:
        aggregate_file.write('Avg ' + heading + '\t StdDev ' + heading + '\t' )
    aggregate_file.write("\n")
    for k in range(0,len(file_lines)-1):
        for col in range(0, len(headings)):
            #print '%s - %s' % (headings[col],values[k][col])
            avg = np.average(values[k][col])
            dev = np.std(values[k][col])
            if col>0:
                aggregate_file.write( str(avg) + '\t' + str(dev)+ '\t')
            else:
                aggregate_file.write(str(avg) + '\t') # for k
        aggregate_file.write("\n")
    aggregate_file.close()


def calculate_class_separation_pr_roc(object_distances, pr_roc_file, ignore_pairs_ground_truth_value=None, object_to_class_func = None, increment_indices = None, verbose=True):
    """
    Calculates Precision/Recall from object pairwise distances with the respect to their classes
    :param object_distances: Sorted list of object distances, where each element in the list is of the form (distance,object_name_1,object_name_2)
    :param pr_roc_file: Output file
    :param object_to_class_func: Function to get object class from object name (if not give, assume that object_distances already contains classes)
    :param increment_indices: Indices for Precision/Recall calculations increments (if None, use 1%,2%,3%,....,100%)
    :return:
    """

    temp_increment_indices = [] #indices in object_distances where goes the next distance increment
    num_pairs = len(object_distances)

    if ignore_pairs_ground_truth_value is not None:
        ignored_pairs_count = 0
        for value in object_distances:
            class_1 = object_to_class_func(value[1]) if object_to_class_func is not None else value[1]
            class_2 = object_to_class_func(value[2]) if object_to_class_func is not None else value[2]
            if ignore_pairs_ground_truth_value is not None and class_1 == class_2 == ignore_pairs_ground_truth_value:
                ignored_pairs_count+=1
        num_pairs -= ignored_pairs_count

    temp_index_increments_from1percent =  list(np.linspace(0,num_pairs,101))
    num_of_splits_upto1percent = 300 if num_pairs>100000 else 100 if num_pairs>10000 else 10 if num_pairs>1000 else 1
    temp_index_increments_to1percent = list(np.linspace(temp_index_increments_from1percent[0],temp_index_increments_from1percent[1],num_of_splits_upto1percent + 1))
    temp_index_increments = temp_index_increments_to1percent + temp_index_increments_from1percent[2:] # up to 1% - small increments, then increments 1%

    k=0
    current_index_increment = 0
    prev_value = -1
    same_class_list = [] #boolean array -> 1 - corresponding pair of objects from the same class, 0 - otherwise
    i=0
    for value in object_distances:
        class_1 = object_to_class_func(value[1]) if object_to_class_func is not None else value[1]
        class_2 = object_to_class_func(value[2]) if object_to_class_func is not None else value[2]
        if ignore_pairs_ground_truth_value is not None and class_1 == class_2 == ignore_pairs_ground_truth_value:
            continue

        same_class_list.append(1 if class_1 == class_2 else 0)

        if i >= current_index_increment and value[0]>prev_value:
            temp_increment_indices.append(i)
            k+=1
            current_index_increment = temp_index_increments[k] if k<len(temp_index_increments)-1 else temp_index_increments[-1]+1
        i+=1
        prev_value = value[0]

    print 'Fraction of same class pairs: %s/%s = %.6f\n' % (sum(same_class_list),(num_pairs),1.0*sum(same_class_list)/(num_pairs))

    for i in range(0, len(temp_index_increments) - len(temp_increment_indices)):
        temp_increment_indices.append(num_pairs)

    if increment_indices is None:
        increment_indices = temp_increment_indices

    #calculate precision/recall
    #for area under the curve
    precisions = []
    recalls = []
    tp_rates = []
    fp_rates = []

    if not os.path.exists(os.path.dirname(pr_roc_file)):
        os.makedirs(os.path.dirname(pr_roc_file))

    with open(pr_roc_file,'wb') as outfile:
        outfile.write("k\tTP\tFP\tTN\tFN\tRecall\tPrecision\tTP rate\tFP rate\tF-score\tAccuracy\n")
        for i in range(0, len(increment_indices)):
            current_index = increment_indices[i]
            k  = 1.0 * current_index / num_pairs
            tp = same_class_list[0:current_index].count(1)
            fp = current_index - tp
            fn = same_class_list[current_index:].count(1)
            tn = num_pairs - current_index - fn

            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)

            outfile.write(str(k) + "\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")

            if verbose:
                print 'k=%.6f: tp=%d, fp=%d, tn=%d, fn=%d, precision=%.3f, recall=%.3f, tp-rate=%.3f, fp-rate=%.3f, f-score=%.3f, accuracy=%.3f' % (k,tp,fp,tn,fn,precision, recall, tp_rate, fp_rate, f_score,accuracy)
            else:
                sys.stdout.write('\rk=%.6f...' % k)
                sys.stdout.flush()

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

    print 'AUPR  = %.3f' % area_under_the_curve(recalls,precisions)
    print 'AUROC = %.3f' % area_under_the_curve(fp_rates,tp_rates)