#!/usr/bin/env python
"""
Frontend for evaluation routines on PARSE-27k dataset.
This can be used to evaluate model predictions against the groundtruth.

The groundtruth is one of the HDF5 files generated by preprocess_dataset.py.
The predictions should be in an HDF5 file - one HDF5 'dataset' for each attribute.
"""

from __future__ import print_function
import sys
from argparse import ArgumentParser
import numpy as np
import h5py
from collections import OrderedDict

import pmetrics

def read_hdf5_groundtruth(filename):
    '''read labels attributes from an HDF5 file
    Returns:
      tuple of (names, labels matrix)
       - names is an array of strings specifying the attribute names
       - labels matrix is an N x K matrix where K is len(names)
    '''
    h5 = h5py.File(filename, 'r')
    out = h5['labels'][:]
    names = h5['labels'].attrs['names']
    h5.close()
    return names, out

def evaluate_retrieval(groundtruth, predictions):
    '''
    Evaluation for the **retrieval viewpoint**
      i.e. the first class (idx 0) is treated as N/A (or ?) class
      For binary (?, yes, no)-type attributes this computes the AP

    Parameters:
      groundtruth: numpy array-like  N (N samples - int label specifying the class)
      predictions: numpy array-like  N x K

    Returns:
      dictionary with several evaluation metrics
    '''
    measures = {}
    confmat = pmetrics.softmax_to_confmat(groundtruth, predictions)

    # check if there were predictions for the first class
    # which is not supposed to happen in retrieval mode
    if np.any(confmat[:,0] > 0):
        raise ValueError('retrieval mode - but there were predictions for N/A class')

    # this compute the two class accuracy - it does not reflect performance on the
    # exmaples with N/A groundtruth label
    measures['accuracy'] = pmetrics.accuracy(confmat, ignore_na=True)

    if predictions.shape[1] == 3:
        # predictions need to be converted to a continous score before
        pred_scores = pmetrics.softmax_prediction_to_binary(predictions)
        # convert groundtruth to [-1, 1] -- groundtruth input is [0,1,...N]
        gt_bin = pmetrics.labels_to_binary(groundtruth)

        measures['AP'] = pmetrics.average_precision_everingham(gt_bin, pred_scores)
    else:
        # non-binary case, for example Orientation, Orientation8, etc.
        # compute specific measures here... if you want.
        pass
    return measures

def evaluate_classification(groundtruth, predictions):
    '''
    Evaluation for the **classification mode**
      i.e. all classes in the groundtruth are 'counted'

    Parameters:
      groundtruth: numpy array-like  N (N samples - int label specifying the class)
      predictions: numpy array-like  N x K

    Returns:
      dictionary with several evaluation metrics
    '''
    measures = {}
    confmat = pmetrics.softmax_to_confmat(groundtruth, predictions)
    measures['accuracy'] = pmetrics.accuracy(confmat)
    measures['BER'] = pmetrics.balanced_error_rate(confmat)
    return measures
#
# -----------------------------------------------------------------------------
#
def evaluate_summary_measures(eval_results, result_field='Summary'):
    '''
    Compute averages and the like over results for several attributes.
    Here, this just computes the mAP and mBER

    Parameters:
      eval_results: dict - attribute -> measures

    Returns:
      the eval_results with added result_field (default: 'Summary')
    '''
    aps = []
    bers = []
    for attr, measures in eval_results.items():
        if 'AP' in measures.keys():
            aps.append(measures['AP'])
        if 'BER' in measures.keys():
            bers.append(measures['BER'])
    if len(aps) > 0:
        eval_results[result_field] = dict(mAP=np.mean(aps))
    if len(bers) > 0:
        eval_results[result_field] = dict(mBER=np.mean(bers))
    return eval_results

def generate_report(eval_results):
    '''
    Parameters:
        eval_results: takes a dict - attribute -> measures
           where measures is a dict - eval_measures -> values
    Returns:
        a string representation of the results
    '''
    out = 80*'='+'\n\t EVALUATION REPORT \n\n'
    for attrib, results in eval_results.items():
        out += 80*'-' + '\n == ' + attrib + ' == \n'
        for name, measure in results.items():
            out += '%15s:    %s\n' % (str(name), str(measure))
        out += '\n'
    return out

# ------------------------------------------------------------------------------
# below this point are the routines for using this module as a standalone tool

def main(args):
    ''' this is run when used directly from the commandline
    '''
    h5_predictions = h5py.File(args.predictions, 'r')
    attributes = h5_predictions.keys()
    names, labels = read_hdf5_groundtruth(args.groundtruth)
    groundtruth = {}
    for idx, name in enumerate(names):
        groundtruth[name] = labels[:, idx]

    eval_results = OrderedDict()
    report = ''
    for attr in attributes:
        if args.verbose:
            print('evaluating attribute: ', attr)
        pred = h5_predictions[attr]
        try:
            gt = groundtruth[attr]
        except KeyError:
            print('no groundtruth for attribute: ', attr)
            sys.exit(2)

        if args.mode == 'retrieval':
            eval_results[attr] = evaluate_retrieval(gt, pred)
        elif args.mode == 'classification':
            eval_results[attr] = evaluate_classification(gt, pred)
        else:
            raise NotImplemented('evaluation mode not implemented.')

    eval_results = evaluate_summary_measures(eval_results)

    report = generate_report(eval_results)
    # output the results
    if args.verbose or args.output is None:
        print(report)
    if args.output:
        f = open(args.output+'.txt', 'w')
        f.write(report)
        f.close()

if __name__ == '__main__':
    parser = ArgumentParser(description=__doc__)
    parser.add_argument('--verbose', '-v', action='store_true')
    parser.add_argument('groundtruth', help='HDF5 file with groundtruth labels')
    parser.add_argument('predictions', help='HDF5 file with model predictions')
    parser.add_argument('--output', '-o', help='filename for eval results report file')
    parser.add_argument('--mode', '-m',
                        choices=('retrieval', 'classification'), default='retrieval',
                        help='evaluate with N/A labels (i.e. BER) '
                             '-- or -- retrieval viewpoint (AP)')
    args = parser.parse_args()
    main(args)
