import re

from transformers import pipeline, AutoTokenizer


## Masked Language Model scoring

# We use the the Roberta model that was pretrained with the TopicalChat data by the USR team
# as a model for gettting the averaged token likelihood of the target sentence.
# The function *sentence_likelihood* also returns the most likeley sentence according to the model


# and the averaged score for the mostly likely tokens.
# We can thus compare the actual response in a turn with the response that would be generated by the LM.

# The LM will return a number of results for the masked token with probability scores.
# We compare the target token with the results to get the score for the target token.
# If the target token is not in the results, we set the probability to "0".

# You can set the number of results returned by the model.
# The more results, the more likely the target gets a score, albeit a very low score.


class MLM:
    def __init__(self, path=None, top_results=20):
        """ Load pretrained RoBERTa model for masked Langauge model based likelihood.

            params
            str path: path to stored model or None

            returns: None
        """
        if path is None:
            self.__model_name = 'google-bert/bert-base-multilingual-cased'
        else:
            self.__model_name = path

        print('Extracting the likelihood score using', self.__model_name)

        self.__tokenizer = AutoTokenizer.from_pretrained(self.__model_name, local_files_only=True)
        self.__model = pipeline("fill-mask", model=self.__model_name)
        self.__model.top_k = top_results  ### we check against the top results

    def mask_target_sentence(self, context, target):
        masked_targets = []
        ## We limit the length of the target as too long utterance break the token limit
        target_tokens = re.split(' ', target[:500])
        for index, token in enumerate(target_tokens):
            sequence = context + " "
            for token in target_tokens[:index]:
                sequence += token + " "
            sequence += self.__tokenizer.mask_token
            for token in target_tokens[index + 1:]:
                sequence += " " + token
            masked_targets.append(sequence)
        return masked_targets, target_tokens

    def sentence_likelihood(self, context, target):
        masked_targets, target_tokens = self.mask_target_sentence(context, target)
        expected_target = ""
        max_scores = []
        scores = []
        for masked_target, token in zip(masked_targets, target_tokens):
            results = self.__model(masked_target)
            expected_target += results[0]['token_str'] + " "
            max_scores.append(results[0]['score'])
            match = False
            for result in results:
                if result['token_str'].lower().strip() == token.lower():
                    scores.append(result['score'])
                    match = True
                    break

            if not match:
                scores.append(0)
        likelihood = sum(scores) / len(scores)
        max_likelihood = sum(max_scores) / len(max_scores)

        return likelihood, expected_target, max_likelihood

    def score_pairs_for_likelihood(self, turns: []):
        for context, target in turns:
            llh, best_sentence, max_score = self.sentence_likelihood(context, target)
            print('Likelihood:', llh, 'Max score:', max_score, 'Best sentence:', best_sentence)


if __name__ == "__main__":
    turns = [('Do you have a cat?', 'I do not have a cat'),  # good
             ('Do you have a cat?', 'I like cats'),  # not as good
             ('Do you have a cat?', 'I like kittens'),  # worse
             ('Do you have a cat?', 'I want a turtle')]  # what are we even saying
    ###### Likelihood
    top_results = 20
    # model_path = 'adamlin/usr-topicalchat-roberta_ft'
    # model_path = 'xlm-roberta-base'
    model_path = 'roberta-base'
    model_mlm = MLM(model_path, top_results)
    for context, target in turns:
        llh, best_sentence, max_score = model_mlm.sentence_likelihood(context, target)
        print(target)
        print('Likelihood:', llh, 'Max score:', max_score, 'Best sentence:', best_sentence)
