"""
LICENCE
-------
Copyright 2013 by Kitware, Inc. All Rights Reserved. Please refer to
KITWARE_LICENSE.TXT for licensing information, or contact General Counsel,
Kitware, Inc., 28 Corporate Drive, Clifton Park, NY 12065.

"""

import os
import numpy as np
import itertools

def compute_errorRate(labels, scores, target_class=1):
    """
    Compute prbability of false alarm (PFA) and probability of miss detection (PD) at every sample point.
    This is done by sorting the scores, and counting the miss/false alarms based on labels.

    @param labels: list of ground truth labels
    @type labels: iterable
    @param scores: scores generated by classifiers
    @type scores: numpy.iterable (includes most iterable)
    @target_class: the label of the target (positive) class (default = 1)
    @type target_class: int
    @return: (sorted) PFAs, PMDs, and sorted indices of the original inputs (increasing order)
    @rtype: tuple of (numpy.array of floats, numpy.array of floats, numpy.array of int)
    """
    num_scores = len(scores)
    if not len(labels)==num_scores:
        print 'The number of labels and scores should be same...'
        return 0
    sortedIndex = np.argsort(scores)[::-1]
    sortedLabels = [labels[idx] for idx in sortedIndex]
#    sortedScores = scores[sortedIndex]

    num_pos = 0
    for i in range(num_scores):
        if labels[i]==target_class:
            num_pos += 1
    num_neg = num_scores - num_pos

    pfas = np.zeros(num_scores)
    pmds = np.zeros(num_scores)
    fp, fn = 0, num_pos

    for i in range(num_scores):
        if sortedLabels[i]==target_class:
            fn -= 1
        else:
            fp += 1
        if not num_neg==0:
            pfas[i] = fp/float(num_neg)
        if not num_pos==0:
            pmds[i] = fn/float(num_pos)
    return pfas, pmds, sortedIndex


def plot_yearly_goals(ax, goal_years):

    intercepts = [0.25, 0.50, 0.65, 0.75, 0.82]

    for idx, year in enumerate(goal_years):
        xs = np.arange(0, 1, 0.0001)
        ys = 12.5 * xs + intercepts[year-1]
        plot_label = 'Year %d Goal' % (year)
        ax.plot(xs*100, ys,'k--', label=plot_label)


def average_precision_R0(inputs, plot_parameters = None):
    """
    Compute Average Precision, R0, and maximum possible R0.
    Also, plots curves if filename is specified.

    @param inputs: list of dictionary, each dictionary contains information regarding different event as follows
        input['scores']: series of scores in the same order as labels, with type iterable
        input['labels']: series of ground truth labels as labels
        input['gamma']: optional, default = 0.001, ratio of positives
        input['target_class']: target class label, integer
        input['threshold']: optional float, if exists, computes R0 (in addition to R0*), and also mark on curve
        input['topN']: optional, if exists, only use topN scores for results computing
        input['plot_parameters']: optional dictionary, if exists each input write its own plot file
            input['plot_parameters']['filename_plot']: string, write a plot file
            input['plot_parameters']['title']: optional, title of image
            input['plot_parameters']['name']: name of this curve to be included in legend
            input['plot_parameters']['color']: optional, color for this curve,
            input['plot_parameters']['linewidth']: optional, linewidth for this curve
            input['plot_parameters']['resolution']: optional, size of output image, default = (600.0, 600.0)
            input['plot_parameters']['goal_years']: optional, list of years to show goal lines. Default = [1,2,3,4,5]
            input['plot_parameters']['use_log_xaxis']: optional boolean, default = False, if True use [10, 100, 1000,..] as uniformly spaced x tickmarks
            input['plot_parameters']['show_threshold']: optional boolean, default = True, if True and thrshold exists, mark the threshold on the curve
            input['plot_parameters']['show_R0_star']: optional boolean, default = True, if True, show R0* in legend

    @param plot_parameters: optional, dictionary, plotting parameters to draw multiple inputs together.
            input['plot_parameters']['filename_plot']: string, path for a plot file.
            input['plot_parameters']['title']: optional, title of image.
            input['plot_parameters']['resolution']: optional, size of output image, default = (600.0, 600.0).
            input['plot_parameters']['goal_years']: optional, list of years to show goal lines. Default = [1,2,3,4,5]. Overrides each input goal years.
            input['plot_parameters']['use_log_xaxis']: optional boolean, default = False, if True use [10, 100, 1000,..] as uniformly spaced x tickmarks.
            input['plot_parameters']['show_threshold']: optional boolean, default = True, if True, mark the thresholds (for those it exists) on the curves
            plot_parameters['show_R0_star']: optional boolean, default = True, if True, show R0* in legend
            For legends/colors/linewidths, use name/color/linewidth of each input plot parameters

    @return: list of dictionary, each with the following items
        outputs: list of output dictionary for each input
        output['ap_prime']: normalized ap
        output['ap']: traditional average precision
        output['R0_star']: maximum possible R0
        output['R0']: optional, R0 at threshold
    """
    #matplotlib.use('Agg') # this step makes it work on machines with GUI backend.
    import matplotlib.pyplot as plt

    colors = itertools.cycle([[1,0,0], "b", "g", "c", "m", "y", "k"])
    outputs = [dict() for k in range(len(inputs))]

    # list of plot handles
    p = [0]*len(inputs)

    # if all the curves are to be plotted on same figure, initialize figure
    if plot_parameters is not None:
        fig = plt.figure()
        ax = fig.add_subplot(1,1,1)


    # iterate over list of inputs
    for idx, ip in enumerate(inputs):

        scores = ip['scores']
        labels = ip['labels']
        num_scores = len(scores)
        if not len(labels)==num_scores:
            print 'The number of labels and scores should be the same.'
            return 0

        sortedIndex = np.argsort(scores)[::-1]
        sortedLabels = [labels[index] for index in sortedIndex]

        # set defaults
        target_class = 1
        gamma = 0.001
        dpi_value = 600
        line_width = 2.0
        show_threshold = True
        goal_years = [1,2,3,4,5]
        show_R0_star = True

        # input parameter options
        if 'target_class' in ip:
            target_class = ip['target_class']

        if 'topN' in ip:
            sortedLabels = sortedLabels[:ip['topN']]
            num_scores = ip['topN']

        if 'gamma' in ip:
            gamma = ip['gamma']


        # compute metrics
        num_pos = sortedLabels.count(target_class)
        percent_rank = np.zeros(num_scores)
        precision = np.zeros(num_scores)
        recall = np.zeros(num_scores)

        true_detection = 0;

        for i in range(num_scores):
            percent_rank[i] =  (i+1)  / float(num_scores)
            if sortedLabels[i] == target_class:
                true_detection += 1

            precision[i] = true_detection / float(i + 1)
            recall[i] = true_detection / float(num_pos)


        indices = [i for i, x in enumerate(sortedLabels) if x == target_class]
        ap_prime = np.sum (gamma * recall[indices] / percent_rank[indices]) / num_pos
        ap = np.mean(precision[indices])
        R0_star = np.max(recall - 12.5 * percent_rank)

        outputs[idx]['ap_prime'] = ap_prime
        outputs[idx]['ap'] = ap
        outputs[idx]['R0_star'] = R0_star


        # plotting parameters

        if 'threshold' in ip:
            index = scores.index(ip['threshold'])
            R0 = recall[index] - 12.5*percent_rank[index]
            outputs[idx]['R0'] = R0
            if 'plot_parameters' is not None:
                if 'show_threshold' in plot_parameters:
                    show_threshold = plot_parameters['show_threshold']
            if 'plot_parameters' in ip:
                if 'show_threshold' in ip['plot_parameters']:
                    show_threshold = ip['plot_parameters']['show_threshold']

        else:
            show_threshold = False

        if plot_parameters is not None:
            if 'show_R0_star' in plot_parameters:
                    show_R0_star = plot_parameters['show_R0_star']
        if 'plot_parameters' in ip:
            if 'show_R0_star' in ip['plot_parameters']:
                    show_R0_star = ip['plot_parameters']['show_R0_star']

        if show_R0_star:
            plot_label = 'AP\' = %0.4f, AP = %0.4f, R0* = %0.4f' % (ap_prime, ap, R0_star)
        else:
            plot_label = 'AP\' = %0.4f, AP = %0.4f' % (ap_prime, ap)


        if 'plot_parameters' in ip:
            if 'color' in ip['plot_parameters']:
                color_name = ip['plot_parameters']['color']
            else:
                color_name = next(colors)
            if 'linewidth' in ip['plot_parameters']:
                line_width = ip['plot_parameters']['linewidth']
            if 'goal_years' in ip['plot_parameters']:
                goal_years = ip['plot_parameters']['goal_years']

            if 'name' in ip['plot_parameters']:
                plot_label = '%s, %s' % (ip['plot_parameters']['name'], plot_label)
        else:
            color_name = next(colors)


        # plot percent_rank vs. recall curve
        if plot_parameters is None:
            fig = plt.figure()
            ax = fig.add_subplot(1,1,1)
        p[idx] = ax.plot(percent_rank*100, recall, linewidth=line_width, color = color_name, label = plot_label)
        if show_threshold:
            index = scores.index(ip['threshold'])
            plt.plot(percent_rank[index], recall[index], '^', markersize = 15, color = color_name)


        # save each plot
        if 'plot_parameters' in ip:
            if 'title' in ip['plot_parameters']:
                plt.title(ip['plot_parameters']['title'])
            if 'resolution' in ip['plot_parameters']:
                dpi_value = ip['plot_parameters']['resolution']
            if 'filename_plot' in ip['plot_parameters']:
                plot_yearly_goals(ax, goal_years)
                handles, labels = ax.get_legend_handles_labels()
                ax.legend(handles[::-1], labels[::-1],loc = 'best', prop = {'size': 10})
                plt.xlabel("Percent Rank")
                plt.ylabel("Recall")
                plt.ylim(0,1)
                plt.xlim(0,5)
                plt.grid(b=True, axis='both', linestyle=':', linewidth=0.3)
                plt.savefig(ip['plot_parameters']['filename_plot'], dpi=dpi_value)


    # save final plot with all the curves
    if plot_parameters is not None:
        if 'title' in plot_parameters:
            plt.title(plot_parameters['title'])
        if 'resolution' in plot_parameters:
            dpi_value = plot_parameters['resolution']
        if 'goal_years' in plot_parameters:
                goal_years = plot_parameters['goal_years']
        if 'filename_plot' in plot_parameters:
            plot_yearly_goals(ax, goal_years)
            handles, labels = ax.get_legend_handles_labels()
            ax.legend(handles[::-1], labels[::-1],loc = 'best', prop = {'size': 10})
            plt.xlabel("Percent Rank")
            plt.ylabel("Recall")
            plt.ylim(0,1)
            plt.xlim(0,5)
            plt.grid(b=True, axis='both', linestyle=':', linewidth=0.3)
            plt.savefig(plot_parameters['filename_plot'], dpi=dpi_value)

    return outputs




def average_precision(scores, labels, target_class=1, thresholds = None, topN = None):
    """
    Compute Average Precision for one class, and return AP.

    @param scores: series of scores in the same order as labels
    @param labels: ground truth labels
    @param target_class: the positive class (among possible labels)
    @param topN: (optional) only compute AP among top N scoring examples
    @type labels: iterable
    @type scores: numpy.array_like
    @type target_class: integer
    @type topN: integer, optional, default = None (all samples will be used)
    @return: average precision
    @rtype: float
    """

    if isinstance(scores[0], list):
        inputs = [dict() for k in range(len(scores))]

        for i in range(len(scores)):

            inputs[i]['scores'] = scores[i]
            inputs[i]['labels'] = labels[i]
            inputs[i]['target_class'] = target_class
            inputs[i]['plot_parameters'] = dict()
            if thresholds is not None:
                inputs[i]['threshold'] = thresholds[i]
            if topN is not None:
                inputs[i]['topN'] = topN

        inputs[0]['plot_parameters']['name'] = 'E006'
        inputs[1]['plot_parameters']['name'] = 'E007'
#        inputs[0]['plot_parameters']['filename_plot'] = 'E006.jpg'
#        inputs[1]['plot_parameters']['filename_plot'] = 'E007.jpg'
#        inputs[0]['plot_parameters']['color'] = [1, 0, 0]
#        inputs[1]['plot_parameters']['color'] = 'm'

    else:
        inputs = dict()
        inputs['scores'] = scores
        inputs['labels'] = labels
        inputs['target_class'] = target_class
        if thresholds is not None:
            inputs['thresholds'] = thresholds
        if topN is not None:
            inputs['topN'] = topN


    plot_parameters = dict()
    plot_parameters['resolution'] = 600
    plot_parameters['title'] = 'Performance Metrics'
    plot_parameters['filename_plot'] = 'performance_metrics.jpg'
    plot_parameters['show_R0_star'] = False
    outputs = average_precision_R0(inputs, plot_parameters)

    #outputs = average_precision_R0(inputs)
    return outputs



if __name__ == "__main__" :

    scores6 = (np.loadtxt('scores6.txt')).tolist()
    labels6 = np.loadtxt('labels6.txt').tolist()

    scores7 = (np.loadtxt('scores7.txt')).tolist()
    labels7 = np.loadtxt('labels7.txt').tolist()

    scores = [scores6, scores7]
    labels = [labels6, labels7]


#     scores = ([5, 4, 3, 2, 1], [1, 2, 3, 4, 5])
#    labels = ([1, 1, 0, 1, 0], [1, 0, 0, 1, 1])
#    target_class = [1, 1]
#    thresholds = [3, 3]
#
    metrics = average_precision(scores, labels, 1)




