# Evaluation routines

from collections import namedtuple, defaultdict
import csv 
import os
import pickle

import numpy as np
from scipy.stats import rankdata

from sklearn.metrics import accuracy_score, precision_recall_fscore_support

import warnings

warnings.filterwarnings('ignore', module='sklearn.metrics')

Evaluation = namedtuple('Evaluation', [
    'name', 
    'features',
    'corpus_train', 
    'corpus_test',
    'lexicon',
    'embedding_name',
    'all_unknown',
    'accuracy', 
    'precision_macro',
    'recall_macro',
    'f1_macro', 
    'precision_weighted', 
    'recall_weighted', 
    'f1_weighted',
    'top5_accuracy',
    'mean_rank',
    'mean_reciprocal_rank',
    'ambiguous_accuracy',
    'unambiguous_accuracy',
    'unknown_accuracy',
    'ambiguous_f1_macro',
    'unambiguous_f1_macro',
    'unknown_f1_macro',
    'baseline_name',
    'not_baseline_accuracy',
    'baseline_accuracy',
    'baseline_agreement',
    'runtime'
])

AveragedEvaluation = namedtuple('AveragedEvaluation', 
    list(Evaluation._fields) + [
    'accuracy_std',
    'precision_macro_std',
    'recall_macro_std',
    'f1_macro_std',
    'precision_weighted_std',
    'recall_weighted_std',
    'f1_weighted_std',
    'top5_accuracy_std',
    'mean_rank_std',
    'mean_reciprocal_rank_std',
    'ambiguous_accuracy_std',
    'unambiguous_accuracy_std',
    'unknown_accuracy_std',
    'ambiguous_f1_macro_std',
    'unambiguous_f1_macro_std',
    'unknown_f1_macro_std',    
    'not_baseline_accuracy_std',
    'baseline_accuracy_std',
    'baseline_agreement_std',
    'runtime_std'
])

def save_evaluations(evaluations, filename, fields, verbose=True):
    os.makedirs(os.path.dirname(filename), exist_ok=True)
    with open(filename,'w') as out:
        csv_out = csv.writer(out, delimiter = '\t')
        csv_out.writerow(list(fields))
                
        for evaluation in evaluations:
            data = ['{0:.3f}'.format(x) if isinstance(x, float) else x for x in evaluation]
            csv_out.writerow(data)

            if verbose: print(data)

class EvaluationManager(object):

    def __init__(self):
        self.evaluations = []

    def write_csv(self, filename, verbose=True):
        save_evaluations(self.evaluations, filename, Evaluation._fields, verbose=verbose)
        file_path, file_ending = os.path.splitext(filename)
        averaged_file_name = file_path + '_averaged' + file_ending        
        save_evaluations(sorted(self.__average_evaluations()), averaged_file_name, AveragedEvaluation._fields, verbose=verbose)

    def add_evaluation(self, name, features, corpus_train, corpus_test, lexicon, embedding_name, all_unknown, y_true, y_pred, lemma_poses, baseline, runtime):
        y_true = np.array(y_true)
        y_pred = np.array(y_pred)
        onehot_prediction = False

        data = {}

        # Remove one-hotness
        if y_true.ndim == 2:
            y_true = np.argmax(y_true, axis=1)

        if y_pred.ndim == 2:
            onehot_prediction = True
            y_pred_onehot = y_pred
            y_pred = np.argmax(y_pred, axis=1)
            data['y_pred_onehot'] = y_pred_onehot

        data['y_true'] = y_true
        data['y_pred'] = y_pred
        data['id_to_frame'] = lexicon.id_to_frame
        data['frame_to_id'] = lexicon.frame_to_id

        # prediction_path = os.path.join('out', 'predictions', '{0}_{1}.pkl'.format(name, all_unknown))
        # with open(prediction_path, 'wb') as f:
        #     pickle.dump(data, f)

        accuracy = accuracy_score(y_true, y_pred)
        precision_macro, recall_macro, f1_macro, _ = precision_recall_fscore_support(y_true, y_pred, average='macro')
        precision_weighted, recall_weighted, f1_weighted, _ = precision_recall_fscore_support(y_true, y_pred, average='weighted')
        top5_accuracy = self.__top5_accuracy(y_true, y_pred_onehot) if onehot_prediction else np.nan
        mean_rank, mean_reciprocal_rank = self.__mean_rank(y_true, y_pred_onehot) if onehot_prediction else (np.nan, np.nan)
        ambiguous_accuracy = self.__ambiguous_accuracy(lexicon, y_true, y_pred, lemma_poses)
        unambiguous_accuracy = self.__unambiguous_accuracy(lexicon, y_true, y_pred, lemma_poses)
        unknown_accuracy = self.__unknown_accuracy(lexicon, y_true, y_pred, lemma_poses)
        ambiguous_f1_macro = self.__ambiguous_f1_macro(lexicon, y_true, y_pred, lemma_poses)
        unambiguous_f1_macro = self.__unambiguous_f1_macro(lexicon, y_true, y_pred, lemma_poses)
        unknown_f1_macro = self.__unknown_f1_macro(lexicon, y_true, y_pred, lemma_poses)        
        not_baseline_accuracy = self.__not_baseline_accuracy(y_true, y_pred, lemma_poses, baseline) if baseline else 0.0
        baseline_accuracy = self.__baseline_accuracy(y_true, y_pred, lemma_poses, baseline) if baseline else 0.0
        baseline_agreement = self.__baseline_agreement(y_pred, lemma_poses, baseline) if baseline else 0.0

        evaluation = Evaluation(name, 
            self.__features_to_name(features),
            corpus_train, 
            corpus_test,
            lexicon.source,
            embedding_name,
            all_unknown, 
            accuracy * 100, 
            precision_macro * 100., recall_macro * 100., f1_macro * 100., 
            precision_weighted * 100., recall_weighted * 100., f1_weighted * 100.,
            top5_accuracy * 100.,
            mean_rank,
            mean_reciprocal_rank,
            ambiguous_accuracy * 100.,
            unambiguous_accuracy * 100.,
            unknown_accuracy * 100.,
            ambiguous_f1_macro * 100.,
            unambiguous_f1_macro * 100.,
            unknown_f1_macro * 100.,            
            type(baseline).__name__,
            not_baseline_accuracy * 100.,
            baseline_accuracy * 100.,
            baseline_agreement * 100.,
            runtime)
        self.evaluations.append(evaluation)

        return evaluation

    def __top5_accuracy(self, y_true, y_pred_onehot):
        # Find indices of the largest 5 elements for each row
        idx = np.argpartition(y_pred_onehot.T, -5, axis=0)[-5:].T

        # Check for every row whether true label is in top 5 results
        scores = np.empty_like(y_true)
        for i, (true_label, top5_results) in enumerate(zip(y_true, idx)):
            scores[i] = true_label in top5_results

        accuracy = np.mean(scores)
        return accuracy

    def __mean_rank(self, y_true, y_pred_onehot):
        ranks = np.empty_like(y_true)
        for i, (true_label, predictions) in enumerate(zip(y_true, y_pred_onehot)):
            # Compute the rank the prediction gave to the true label
            # Negative because rankdata normally ranks from smallest to largest
            rank = rankdata(-predictions)[true_label]
            ranks[i] = rank
        return np.mean(ranks), np.mean(1 / ranks)

    def __ambiguous_accuracy(self, lexicon, y_true, y_pred, lemma_poses):
        return self.__filtered_accuracy(lexicon.is_ambiguous, y_true, y_pred, lemma_poses)

    def __unambiguous_accuracy(self, lexicon, y_true, y_pred, lemma_poses):
        return self.__filtered_accuracy(lexicon.is_unambiguous, y_true, y_pred, lemma_poses)

    def __unknown_accuracy(self, lexicon, y_true, y_pred, lemma_poses):
        return self.__filtered_accuracy(lexicon.is_unknown, y_true, y_pred, lemma_poses)

    def __ambiguous_f1_macro(self, lexicon, y_true, y_pred, lemma_poses):
        return self.__filtered_f1_macro(lexicon.is_ambiguous, y_true, y_pred, lemma_poses)

    def __unambiguous_f1_macro(self, lexicon, y_true, y_pred, lemma_poses):
        return self.__filtered_f1_macro(lexicon.is_unambiguous, y_true, y_pred, lemma_poses)

    def __unknown_f1_macro(self, lexicon, y_true, y_pred, lemma_poses):
        return self.__filtered_f1_macro(lexicon.is_unknown, y_true, y_pred, lemma_poses)        

    def __filtered_accuracy(self, f, y_true, y_pred, lemma_poses):
        filtered_idx = [f(lemmapos) for lemmapos in lemma_poses]
        y_true_filtered = y_true[filtered_idx]
        y_pred_filtered = y_pred[filtered_idx]
        return accuracy_score(y_true_filtered, y_pred_filtered)

    def __filtered_f1_macro(self, f, y_true, y_pred, lemma_poses):
        filtered_idx = [f(lemmapos) for lemmapos in lemma_poses]
        y_true_filtered = y_true[filtered_idx]
        y_pred_filtered = y_pred[filtered_idx]
        precision_macro, recall_macro, f1_macro, _ = precision_recall_fscore_support(y_true_filtered, y_pred_filtered, average='macro')        
        return f1_macro

    def __not_baseline_accuracy(self, y_true, y_pred, lemma_poses, baseline):
        baseline_predictions = baseline.predict(lemma_poses)
        different_idx = [yt != yb if yb >= 0 else False for yt, yb in zip(y_true, baseline_predictions)]
        y_true_different = y_true[different_idx]
        y_pred_different = y_pred[different_idx]
        return accuracy_score(y_true_different, y_pred_different)

    def __baseline_accuracy(self, y_true, y_pred, lemma_poses, baseline):
        baseline_predictions = baseline.predict(lemma_poses)
        same_idx = [yt == yb if yb >= 0 else False for yt, yb in zip(y_true, baseline_predictions)]
        y_true_same = y_true[same_idx]
        y_pred_same = y_pred[same_idx]
        return accuracy_score(y_true_same, y_pred_same)

    def __baseline_agreement(self, y_pred, lemma_poses, baseline):
        baseline_predictions = baseline.predict(lemma_poses)
        same_predictions = [yp == yb if yb >= 0 else False for yp, yb in zip(y_pred, baseline_predictions)]
        return np.mean(same_predictions)

    def __features_to_name(self, features):
        return ', '.join(sorted([feature.name for feature in (features)]))  

    def __average_evaluations(self):
        groups = defaultdict(list)

        # Group by classifier
        for ev in self.evaluations:
            key = (ev.name, ev.features, ev.corpus_train, ev.corpus_test, ev.lexicon, ev.embedding_name, ev.all_unknown, ev.baseline_name)
            groups[key].append(ev)

        averaged_evaluations = []
        for key, group in groups.items():
            averaged_evaluation = self.__average_group(key, group)
            averaged_evaluations.append(averaged_evaluation)  

        return averaged_evaluations      

    def __average_group(self, key, group):
        common_fields = ['name', 'features', 'corpus_train', 'corpus_test', 'lexicon', 'embedding_name', 'all_unknown', 'baseline_name']    
        averaged_fields = ['accuracy', 'precision_macro', 'recall_macro', 'f1_macro', 'precision_weighted', 'recall_weighted', 'f1_weighted',
         'top5_accuracy', 'mean_rank', 'mean_reciprocal_rank', 'ambiguous_accuracy', 'unambiguous_accuracy', 'unknown_accuracy', 'ambiguous_f1_macro',
         'unambiguous_f1_macro', 'unknown_f1_macro', 'not_baseline_accuracy', 'baseline_accuracy', 'baseline_agreement', 'runtime']

        values = defaultdict(list)

        for evaluation in group:
            for field in averaged_fields:
                value = getattr(evaluation, field)
                values[field].append(float(value))

        means = {k : np.mean(v) for k,v in values.items()}
        stds = {k + '_std' : np.std(v) for k,v in values.items()}

        merged_fields = dict(zip(common_fields, key))
        merged_fields.update(means)
        merged_fields.update(stds)

        return AveragedEvaluation(**merged_fields)    

def predict_with_lexicon(model, X, y_true, lemma_pos, lexicon):
    predictions = model.predict(X)
    predictions_with_lexicon = []

    for lemmapos, onehot_prediction in zip(lemma_pos, predictions):
        predicted_frame = predict_with_lexicon_single(lexicon, lemmapos, onehot_prediction)

        predictions_with_lexicon.append(predicted_frame)
    
    return np.array(predictions_with_lexicon)

def predict_with_lexicon_single(lexicon, lemmapos, onehot_prediction):
    if lexicon.is_unknown(lemmapos):
        available_frame_ids = lexicon.get_all_frame_ids()    
    else:            
        available_frame_ids = lexicon.get_available_frame_ids(lemmapos)

    if len(available_frame_ids) == 1:
        # if the LU is known and has only one frame, just return it. Even if there is no data for this LU (!)
        return available_frame_ids[0]
    else:
        # pick the best-scoring frame among available
        available_predictions = onehot_prediction[available_frame_ids]
        best_class_idx = np.argmax(available_predictions)
        return available_frame_ids[best_class_idx]    
