from collections import defaultdict, namedtuple
import csv
from enum import Enum, unique
import glob
import logging

import numpy as np

from keras.models import load_model

import h5py

from sklearn.metrics import confusion_matrix

from basic_srl.evaluation import EvaluationManager, predict_with_lexicon
from basic_srl.formats import load_annotations
from basic_srl.globals import *
from basic_srl.resources import ResourceManager
from basic_srl.lexicon import load_lexicon
from basic_srl.mlp_classifier import MlpFeature, load_dataset_mlp

ModelMetadata = namedtuple('ModelMetadata', [
    'corpus_train',
    'corpus_validation',
    'embedding_name'
])

SEQUENCE_LENTH = 60

@unique
class ModelType(Enum):
    MLP = 1

def write_csv(output_file, header, entries):
    with open(output_file,'w') as f:
        csv_out = csv.writer(f, delimiter = '\t')
        csv_out.writerow(list(header))

        for entry in entries:
            data = ['{0:.2f}'.format(x) if isinstance(x, float) else x for x in entry]
            csv_out.writerow(data)
    logging.info('Written [%d] entries to [%s]', len(entries), output_file)    

def remove_onehot(x):
    return np.argmax(x, axis=1) if x.ndim == 2 else x

def get_best_model(resource_manager, model_base_name, model_type, corpus_name, lexicon_name, embeddings):
    model_pattern = resource_manager.get_model(model_base_name + '-*.h5')
    model_files = glob.glob(model_pattern)
    evaluation_manager = EvaluationManager()

    highest_f1_macro = 0
    best_model_name = None

    for model_name in model_files:
        evaluation, evaluation_with_lexicon = evaluate_model(resource_manager, evaluation_manager, model_name, model_type, corpus_name, lexicon_name, embeddings)

        if evaluation_with_lexicon.f1_macro > highest_f1_macro:
            highest_f1_macro = evaluation_with_lexicon.f1_macro
            best_model_name = model_name

    return best_model_name

def load_trained_model(model_name, model_type):
    if model_type == ModelType.MLP:
        feature_enum = MlpFeature
    else:
        raise Exception('Unknown model type: ' + str(model_type))    

    frame_to_id_mapping = {}
    features = []

    with h5py.File(model_name, 'r') as f:
        grp = f['basic_srl']
        embedding_name = grp.attrs['embedding_name']
        corpus_train = grp.attrs['corpus_train']
        corpus_val = grp.attrs['corpus_validation']

        # Read frame-to-id-mapping
        for i, frame in enumerate(grp['frame_ids']):
            frame_to_id_mapping[frame] = i

        # Read features
        for feature_id in grp['features']:
            features.append(feature_enum(feature_id))
    
    logging.info('Loaded [%s] with features [%s] and embeddings [%s]', model_name, features, embedding_name)

    metadata = ModelMetadata(corpus_train, corpus_val, embedding_name)
    return load_model(model_name), frame_to_id_mapping, set(features), metadata

def load_dataset(resource_manager, model_type, corpus_name, embeddings, lexicon, features):
    if model_type == ModelType.MLP:
        
        if isinstance(embeddings, dict): # embeddings are given as a dict
            dict_embeddings = embeddings
        else: # embeddings are given directly as one textual embedding instance
            # make a dict out of it:
            dict_embeddings = {}
            dict_embeddings["embeddings"] = embeddings    

        return load_dataset_mlp(resource_manager, corpus_name, dict_embeddings, lexicon, features)
    else:
        raise Exception('Unknown model type: ' + str(model_type))

def evaluate_model(resource_manager, evaluation_manager, model_name, model_type, corpus_name, lexicon_name, embeddings):
    model, frame_to_id_mapping, features, metadata = load_trained_model(model_name, model_type)
    assert metadata.embedding_name == embeddings.vsm_name, 'Cannot use two different embeddings for train and eval'    

    lexicon = load_lexicon(resource_manager, lexicon_name, frame_to_id_mapping)


    dataset, X, Y = load_dataset(resource_manager, model_type, corpus_name, embeddings, lexicon, features)
        
    predictions = model.predict(X)
    predictions_with_lexicon = predict_with_lexicon(model, X, Y, dataset.lemma_pos, lexicon)

    evaluation = evaluation_manager.add_evaluation(model_name, features, metadata.corpus_train, corpus_name, lexicon, embeddings.vsm_name, True, Y, predictions, dataset.lemma_pos, None, 0)
    evaluation_with_lexicon = evaluation_manager.add_evaluation(model_name, features, metadata.corpus_train, corpus_name, lexicon, embeddings.vsm_name, False, Y, predictions_with_lexicon, dataset.lemma_pos, None, 0)            

    return evaluation, evaluation_with_lexicon    

def error_analysis(resource_manager):
    def analyse_errors(model_name, model_type, corpus_test, lexicon_name, embeddings, corpus_traindev = None):       
        Entry = namedtuple('Entry', ['name', 'correct', 'unseen_at_training', 'seen_only_with_different_label', 'normal_error',
            'correct_lex', 'unseen_at_training_lex', 'seen_only_with_different_label_lex', 'normal_error_lex', 'wrong_by_lexicon'])

        model, frame_to_id_mapping, features, metadata = load_trained_model(model_name, model_type)
        
        lexicon = load_lexicon(resource_manager, lexicon_name, frame_to_id_mapping)  

        if isinstance(embeddings, dict): # embeddings are given as a dict
            dict_embeddings = embeddings
            embeddings_vsm_name_concatenatedString = ""
            if "embeddings" in dict_embeddings.keys():
                embeddings_vsm_name_concatenatedString += dict_embeddings["embeddings"].vsm_name
            if "synset_embeddings" in dict_embeddings.keys():
                embeddings_vsm_name_concatenatedString += dict_embeddings["synset_embeddings"].vsm_name
            if "imagined_embeddings" in dict_embeddings.keys():
                embeddings_vsm_name_concatenatedString += dict_embeddings["imagined_embeddings"].vsm_name
            embeddings_vsm_name = embeddings_vsm_name_concatenatedString
            
            dataset_traindev, x_traindev, y_traindev = load_dataset(resource_manager, model_type, corpus_traindev, embeddings, lexicon, features)
            dataset_test, x_test, y_test = load_dataset(resource_manager, model_type, corpus_test, embeddings, lexicon, features)
            
            train_annotations = defaultdict(set)
            test_annotations = defaultdict(set)
            test_labels = []
            
            y_traindev = remove_onehot(y_traindev)
            for y_traindev_i, lemma_pos in zip(y_traindev, dataset_traindev.lemma_pos):
                frame_name = lexicon.get_frame(y_traindev_i)
                train_annotations[frame_name].add(lemma_pos)
                
            y_test = remove_onehot(y_test)
            for y_test_i, lemma_pos in zip(y_test, dataset_test.lemma_pos):
                frame_name = lexicon.get_frame(y_test_i)
                test_annotations[frame_name].add(lemma_pos)
                test_labels.append(lemma_pos)
        
        else: # embeddings are given directly as one textual embedding instance
            # make a dict out of it:
            dict_embeddings = {}
            dict_embeddings["embeddings"] = embeddings    
            embeddings = embeddings
            embeddings_vsm_name = embeddings.vsm_name
            
            dataset_test, x_test, y_test = load_dataset(resource_manager, model_type, corpus_test, embeddings, lexicon, features)
            
            corpus_train = metadata.corpus_train
            train_annotations_file = resource_manager.get_frame_annotations_file(corpus_train)
            test_annotations_file = resource_manager.get_frame_annotations_file(corpus_test)
    
            train_annotations = defaultdict(set)
            test_annotations = defaultdict(set)
            test_labels = []
    
            for annotation in load_annotations(train_annotations_file):
                train_annotations[annotation.frame_name].add(annotation.lemma_pos)
            
            for annotation in load_annotations(test_annotations_file):
                test_annotations[annotation.frame_name].add(annotation.lemma_pos)
                test_labels.append(annotation.lemma_pos)
        
        assert metadata.embedding_name == embeddings_vsm_name, 'Cannot use two different embeddings for train and eval'    # embeddings.vsm_name, 'Cannot use two different embeddings for train and eval'   

        
        y_pred = model.predict(x_test)
        y_pred_lex = predict_with_lexicon(model, x_test, y_test, dataset_test.lemma_pos, lexicon)

        y_pred = remove_onehot(y_pred)
        y_pred_lex = remove_onehot(y_pred_lex)
        y_true = remove_onehot(y_test)

        n = len(test_labels)
        unseen_at_training = 0
        seen_only_with_different_label = 0
        normal_error = 0
        correct = 0

        unseen_at_training_lex = 0
        seen_only_with_different_label_lex = 0
        normal_error_lex = 0
        correct_lex = 0        
        wrong_by_lexicon = 0

        for i in range(n): # range(len(y_pred)): # 
            prediction = y_pred[i]
            prediction_lex = y_pred_lex[i]
            goldlabel = y_true[i]
#             lemma_pos = test_labels[i]
            lemma_pos = dataset_test.lemma_pos[i]

            predicted_frame = lexicon.get_frame(prediction)
            gold_frame = lexicon.get_frame(goldlabel)

            # Without lexicon 

            if prediction == goldlabel:
                correct += 1
            # No data in train but in test
            elif gold_frame not in train_annotations:
                assert gold_frame in test_annotations  
                unseen_at_training += 1
            # Different label in train than in test                
            elif gold_frame in train_annotations and lemma_pos not in train_annotations[gold_frame]:
                assert gold_frame in test_annotations                  
                seen_only_with_different_label += 1
            else:
                normal_error += 1

            # With lexicon
            if prediction_lex == goldlabel:            
                correct_lex += 1
            # Lexicon contains the lemma pos but not the right label
            elif goldlabel not in lexicon.get_available_frame_ids(lemma_pos):
                wrong_by_lexicon += 1
            # No data in train but in test
            elif gold_frame not in train_annotations:
                assert gold_frame in test_annotations  
                unseen_at_training_lex += 1
            # Different label in train than in test                
            elif gold_frame in train_annotations and lemma_pos not in train_annotations[gold_frame]:
                assert gold_frame in test_annotations                  
                seen_only_with_different_label_lex += 1
            else:
                normal_error_lex += 1
            
        logging.info('Correct: %.2f, Unseen: %.2f, Seen with different label: %.2f, Normal error: %.2f', correct * 100 / n, unseen_at_training * 100 / n, seen_only_with_different_label * 100 / n, normal_error * 100 / n)
        logging.info('Correct: %.2f, Wrong by lexicon: %.2f, Unseen: %.2f, Seen with different label: %.2f, Normal error: %.2f', correct_lex * 100 / n, wrong_by_lexicon * 100 / n, unseen_at_training_lex * 100 / n, seen_only_with_different_label_lex * 100 / n, normal_error_lex* 100 / n)

        return Entry(model_name, correct  * 100 / n, unseen_at_training  * 100 / n, seen_only_with_different_label  * 100 / n, normal_error * 100 / n,
            correct_lex  * 100 / n, unseen_at_training_lex  * 100 / n, seen_only_with_different_label_lex  * 100 / n, normal_error_lex * 100 / n, wrong_by_lexicon  * 100 / n)

    # best_model_name_mlp_das = get_best_model(resource_manager, 'das-best-mlp', CORPUS_HERMANN.test, LEXICON_FN15_FULL_SELF, EMBEDDINGS_GLOVE_EN_840B_300d)
    best_model_name_mlp_das = 'out/models/das-best-mlp-5.h5'

    # best_model_name_mlp_salsa = get_best_model(resource_manager, 'salsa-best-mlp', CORPUS_SALSA2_DE.test, LEXICON_SALSA2, EMBEDDINGS_GERM_EVAL_DEPS_MIN_5)
    best_model_name_mlp_salsa = 'out/models/salsa-best-mlp-7.h5'
    
    best_model_name_mlp_das_paper = 'out/models/das-best-mlp-9.h5'
    best_model_name_mlp_mm_das_paper = 'out/models/das-best-mm-im-syns-mlp_smallIKRL-9.h5'
    best_model_name_mlp_salsa_paper = 'out/models/salsa-best-mlp-9.h5'
    best_model_name_mlp_mm_salsa_paper = 'out/models/salsa-best-mm-im-mlp_smallIKRL-7.h5'
    
    
    dict_embeddings_das = {}
    dict_embeddings_das["embeddings"] = EMBEDDINGS_GLOVE_EN_6B_300d
    dict_embeddings_das["synset_embeddings"] = EMBEDDINGS_IKRL_NORM
    dict_embeddings_das["imagined_embeddings"] = EMBEDDINGS_IMAGINED_VANILIA_GLOVE_IKRL
    dict_embeddings_salsa = {}
    dict_embeddings_salsa["embeddings"] = EMBEDDINGS_GERM_EVAL_DEPS_MIN_5
    dict_embeddings_salsa["synset_embeddings"] = EMBEDDINGS_IKRL_NORM
    dict_embeddings_salsa["imagined_embeddings"] = EMBEDDINGS_IMAGINED_VANILIA_GERMAN_5MIN_IKRL
        
    entries = [
#         analyse_errors(best_model_name_mlp_das, ModelType.MLP, CORPUS_HERMANN.test, LEXICON_FN15_FULL_SELF, EMBEDDINGS_GLOVE_EN_6B_300d),
#         analyse_errors(best_model_name_mlp_das, ModelType.MLP, CORPUS_HERMANN.test, LEXICON_FN15_FULL_SELF, EMBEDDINGS_GLOVE_EN_840B_300d),
#         analyse_errors(best_model_name_mlp_salsa, ModelType.MLP, CORPUS_SALSA2_DE.test, LEXICON_SALSA2, EMBEDDINGS_GERM_EVAL_DEPS_MIN_5),

        analyse_errors(best_model_name_mlp_das_paper, ModelType.MLP, CORPUS_HERMANN.test, LEXICON_FN15_FULL_SELF, EMBEDDINGS_GLOVE_EN_6B_300d),
        analyse_errors(best_model_name_mlp_mm_das_paper, ModelType.MLP, CORPUS_HERMANN.test, LEXICON_FN15_FULL_SELF, dict_embeddings_das, corpus_traindev = CORPUS_HERMANN.traindev),
        analyse_errors(best_model_name_mlp_salsa_paper, ModelType.MLP, CORPUS_SALSA2_DE.test, LEXICON_SALSA2, EMBEDDINGS_GERM_EVAL_DEPS_MIN_5),
        analyse_errors(best_model_name_mlp_mm_salsa_paper, ModelType.MLP, CORPUS_SALSA2_DE.test, LEXICON_SALSA2, dict_embeddings_salsa, corpus_traindev = CORPUS_SALSA2_DE.traindev)


    ]

    output_file = resource_manager.get_statistics('error-analysis.csv')
    field_names = ['Name', 'C', 'U', 'UL', 'NE', 'C', 'U', 'UL', 'NE', 'WL']
    write_csv(output_file, field_names, entries)

def build_confusion_matrix(model_name, model_type, corpus_test, lexicon_name, embeddings, file_name):
    def save_confusion_matrix(cnf_matrix, lexicon, file_name, list_labels):
        rows, columns = cnf_matrix.shape
        assert rows == columns
#         n = lexicon.get_number_of_labels()
        n = len(list_labels)
        
        data = [['' for i in range(n+1)] for j in range(n+1)] # i = 0, ..., 895

        for idx in range(n): # idx = 0, ..., 894
#             label = lexicon.get_frame(idx)
            label = lexicon.get_frame(list_labels[idx])            
            data[0][idx+1] = label
            data[idx+1][0] = label

        for row in range(rows): # row = 0, ..., 531 / 515
            for col in range(columns):
                e = cnf_matrix[row][col]
                if e != 0:
                    data[row+1][col+1] = str(e)

        # Write data
        with open(resource_manager.get_statistics(file_name), 'w') as f:
            for idx in range(n+1):
                entry = '\t'.join(data[idx])
                f.write(entry)
                f.write('\n')

    model, frame_to_id_mapping, features, metadata = load_trained_model(model_name, model_type)
    
    if isinstance(embeddings, dict): # embeddings are given as a dict
        dict_embeddings = embeddings
        embeddings_vsm_name_concatenatedString = ""
        if "embeddings" in dict_embeddings.keys():
            embeddings_vsm_name_concatenatedString += dict_embeddings["embeddings"].vsm_name
        if "synset_embeddings" in dict_embeddings.keys():
            embeddings_vsm_name_concatenatedString += dict_embeddings["synset_embeddings"].vsm_name
        if "imagined_embeddings" in dict_embeddings.keys():
            embeddings_vsm_name_concatenatedString += dict_embeddings["imagined_embeddings"].vsm_name
        embeddings_vsm_name = embeddings_vsm_name_concatenatedString
    else: # embeddings are given directly as one textual embedding instance
        # make a dict out of it:
        dict_embeddings = {}
        dict_embeddings["embeddings"] = embeddings    
        embeddings = embeddings
        embeddings_vsm_name = embeddings.vsm_name
    
    assert metadata.embedding_name == embeddings_vsm_name, 'Cannot use two different embeddings for train and eval'   

    corpus_train = metadata.corpus_train

    lexicon = load_lexicon(resource_manager, lexicon_name, frame_to_id_mapping)  
    
    dataset_test, x_test, y_test = load_dataset(resource_manager, model_type, corpus_test, embeddings, lexicon, features)

    y_pred = model.predict(x_test)
    y_pred_lex = predict_with_lexicon(model, x_test, y_test, dataset_test.lemma_pos, lexicon)    

    labels = [lexicon.get_frame(i) for i in range(lexicon.get_number_of_labels())]

    y_true = remove_onehot(y_test)    
    y_pred = remove_onehot(y_pred)
    y_pred_lex = remove_onehot(y_pred_lex)
    
#     y_true = y_true[0:3]    
#     y_pred = y_pred[0:3]
#     y_pred_lex = y_pred_lex[0:3]
    
    j = 0
    labels_pred = []
    labels_pred_lex = []
    for y_true_i, y_pred_i, y_pred_lex_i in zip(y_true, y_pred, y_pred_lex):
#         if lexicon.get_frame(y_true_i) == "Performers_and_roles" or lexicon.get_frame(y_true_i) == "Competition":
#         if lexicon.get_frame(y_true_i) == "Ride_vehicle":
        if lexicon.get_frame(y_true_i) == "Statement":       
            j += 1
            print(j, lexicon.get_frame(y_true_i), lexicon.get_frame(y_pred_i), lexicon.get_frame(y_pred_lex_i))
        labels_pred.append(y_true_i)
        labels_pred.append(y_pred_i)
        labels_pred_lex.append(y_true_i)
        labels_pred_lex.append(y_pred_lex_i)
    list_labels_pred = list(set(labels_pred))
    labels_pred = np.asarray(list_labels_pred)
    list_labels_pred_lex = list(set(labels_pred_lex))
    labels_pred_lex = np.asarray(list_labels_pred_lex)
    
    confusion_matrix_nolex = confusion_matrix(y_true, y_pred, labels=labels_pred)
    confusion_matrix_lex = confusion_matrix(y_true, y_pred_lex, labels=labels_pred_lex)

    save_confusion_matrix(confusion_matrix_nolex, lexicon, '{0}-confusion-nolex.csv'.format(file_name), list_labels_pred)
    save_confusion_matrix(confusion_matrix_lex, lexicon, '{0}-confusion-lex.csv'.format(file_name), list_labels_pred_lex)

if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')

    resource_manager = ResourceManager(HOME)
    error_analysis(resource_manager)

# #     build_confusion_matrix('out/models/das-best-mlp-5.h5', ModelType.MLP, CORPUS_HERMANN.test, LEXICON_FN15_FULL_SELF, EMBEDDINGS_GLOVE_EN_840B_300d, 'das-best')
# 
#     build_confusion_matrix('out/models/das-best-mlp-9.h5', ModelType.MLP, CORPUS_HERMANN.test, LEXICON_FN15_FULL_SELF, EMBEDDINGS_GLOVE_EN_6B_300d, 'das-best')
# #     build_confusion_matrix('out/models/das-best-mlp-5.h5', ModelType.MLP, CORPUS_HERMANN.test, LEXICON_FN15_FULL_SELF, EMBEDDINGS_GLOVE_EN_6B_300d, 'das-best')
#     
#     dict_embeddings = {}
#     dict_embeddings["embeddings"] = EMBEDDINGS_GLOVE_EN_6B_300d
#     dict_embeddings["synset_embeddings"] = EMBEDDINGS_WN_GNEWS_SYNSET_300d
#     dict_embeddings["imagined_embeddings"] = EMBEDDINGS_IMAGINED_VANILIA_VGG_GLOVE_IKRL_large
#     build_confusion_matrix('out/models/das-best-mm-im-mlp-3.h5', ModelType.MLP, CORPUS_HERMANN.test, LEXICON_FN15_FULL_SELF, dict_embeddings, 'das-best-mm-im')
# #     build_confusion_matrix('out/models/das-best-mm-im-mlp-1.h5', ModelType.MLP, CORPUS_HERMANN.test, LEXICON_FN15_FULL_SELF, dict_embeddings, 'das-best-mm-im')
#     dict_embeddings["synset_embeddings"] = EMBEDDINGS_VGG_128d_NORM_IKRL_NORM
#     build_confusion_matrix('out/models/das-best-mm-im-mlp_synsetSize_V-5.h5', ModelType.MLP, CORPUS_HERMANN.test, LEXICON_FN15_FULL_SELF, dict_embeddings, 'das-best-mm-im_synsetSize_V')
# #     build_confusion_matrix('out/models/das-best-mm-im-mlp_synsetSize_V-3.h5', ModelType.MLP, CORPUS_HERMANN.test, LEXICON_FN15_FULL_SELF, dict_embeddings, 'das-best-mm-im_synsetSize_V')
    
    
