import os

import torch
from mosestokenizer import MosesDetokenizer
import pandas as pd
import tqdm

from metrics.collection.MetricClass import MetricClass
import dill

from project_root import ROOT_DIR


class MoverScore(MetricClass):
    '''
    A wrapper class for MoverScore(https://github.com/AIPHES/emnlp19-moverscore) by:
    Wei Zhao, Maxime Peyrard, Fei Liu, Yang Gao, Christian M. Meyer, and Steffen Eger. “MoverScore: Text
    Generation Evaluating with Contextualized Embeddings and Earth Mover Distance”. In: Proceedings
    of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International
    Joint Conference on Natural Language Processing (EMNLP-IJCNLP). Hong Kong, China: Association
    for Computational Linguistics, Nov. 2019, pp. 563–578. doi: 10.18653/v1/D19-1053. url:
    https://www.aclweb.org/anthology/D19-1053.
    '''
    ref_based = True
    name = 'MOVERSCORE'


    def __init__(self, bs=64):
        from metrics.collection.metrics_libs.moverscore.moverscore import get_idf_dict, MoverScoreBase
        self.bs = bs
        self.tokenizer = MosesDetokenizer(lang='en')
        self.ms = MoverScoreBase()
        self.get_idf = lambda lst: get_idf_dict(lst)

    def __call__(self, ref, hyp, idf_dict_ref, idf_dict_hyp, n=2):
        return self.ms.word_mover_score(ref, hyp, idf_dict_ref, idf_dict_hyp, \
                                stop_words=[], n_gram=n, batch_size=self.bs)

    def __del__(self):
        # Just deleting this wrapper will not free gpu
        del self.ms.model
        del self.model
        torch.cuda.empty_cache()

    def get_abstraction(self, ref, hyp_dict, ref_dict):
        return lambda hyp: self.__call__([ref] * len(hyp), hyp, ref_dict, hyp_dict)

    def preprocess_df(self, df, hyp_column='HYP'):
        # Tokenize
        tokenized_input = df.copy()
        tokenized_input['REF'] = df.apply(lambda x: self.tokenizer(x['REF'].split(' ')), axis=1)
        tokenized_input[hyp_column] = df.apply(lambda x: self.tokenizer(x[hyp_column].split(' ')), axis=1)
        return tokenized_input

    def gen_idf_from_corpus(self, pandas_corpus, out_path='wmt_17_idf.dill', hyp_column='HYP', system_column='SYSTEM'):
        '''
        This function generated a s dictionary of idf weights for a given pandas corpus
        :param pandas_corpus: A path to a Pandas Corpus generated by WMT loader
        :param out_path: where to save the precomputet dict
        :param hyp_column: The column name with the hypothesis. E.g. useful for kendall ranking, where the dataframe
                           columns have different names.
        :param system_column: The column with the system name. As the original moverscore evaluation script generates
                              one idf per system. Again this is interesting for the kendall ranking
        :return: {lp1: {ref: ref_idf}, {hyp: {system1 : hyp_idf, system2, ...}
        '''
        print('Loading idfs from wmt')

        if type(pandas_corpus) == str:
            df = pd.read_csv(pandas_corpus, delimiter='\t')
        else:
            df = pandas_corpus

        df = self.preprocess_df(df, hyp_column=hyp_column)

        # get all system names per lp
        lp_list = df['LP'].unique().tolist()
        lp_system_dict = {lp: df[df['LP'] == lp][system_column].unique().tolist() for lp in
                          lp_list}

        # generate ref idfs for each lp and hyp idfs for each lp and system
        idfs = {}
        for lp in tqdm.tqdm(lp_system_dict, desc='Generating idfs'):
            filteref_df = df[df['LP'] == lp][['REF', hyp_column, system_column]]
            refs = filteref_df['REF'].tolist()
            idfs[lp] = {'ref_idf': self.get_idf(refs), 'hyp_idfs': {}}
            print(idfs)
            for sys in lp_system_dict[lp]:
                hyps = filteref_df[filteref_df[system_column] == sys][hyp_column].tolist()
                idfs[lp]['hyp_idfs'][sys] = self.get_idf(hyps)

        print('Successfully loaded idfs from WMT, pickling ...')
        with open(out_path, 'wb') as pickle_file:
            # Dill provides more options than pickle
            dill.dump(idfs, pickle_file, -1)
        print('Successfully saved idfs to file:', out_path)

        with open(out_path, 'rb') as in_strm:
            new_idfs = dill.load(in_strm)

        return new_idfs

    def load_idf_dict_from_file(self, in_path='wmt_17_idf.dill'):
        with open(in_path, 'rb') as in_strm:
            idfs = dill.load(in_strm)
        return idfs


if __name__ == '__main__':
    b = MoverScore()

    # Sample using ref and hyp lists
    hyp_dict = b.get_idf(["A simple  for test"])
    ref_dict = b.get_idf(["A simple sentence for test"])
    print(b(["A simple  for test"], ["A simple sentence for test"], hyp_dict, ref_dict))
    #[1.0]

    # Sample using a fixed reference for a list of hypothesis
    hyp_dict = b.get_idf(["A test sentence for.", "A test sentence for.", "A test sentence for."])
    b_trimmed = b.get_abstraction("A test sentence for.", hyp_dict, ref_dict)
    print(
        b_trimmed(["A simple sentence for test", "Another simple sentence for test", 'A test sentence for']))
    #[4.679604320156905e-06, 2.876247215177763e-06, 1.0]

    # Samples using precomputed idfs from wmt17 (needs to be precomputed first)
    with open(os.path.join(ROOT_DIR,'metrics/collection/metrics_libs/moverscore/idfs/wmt17_idf.dill'), 'rb') as in_strm:
        idf_dicts = dill.load(in_strm)

    hyp_dict = idf_dicts['de-en']['ref_idf']
    ref_dict = idf_dicts['de-en']['ref_idf']
    print(b(["A simple  for test"], ["A simple sentence for test"], hyp_dict, ref_dict))
    #[0.6388454538956146]

    b_trimmed = b.get_abstraction("A test sentence for.", hyp_dict, ref_dict)
    print(
        b_trimmed(["A simple sentence for test", "Another simple sentence for test", 'A test sentence for']))
    # [0.39384467140358304, 0.32599533476070297, 0.8983243238543307]