"""
Computes the BLEU, ROUGE
using the COCO image caption competition metrics scripts
(Evaluation codes for Microsoft COCO caption generation)
Clone from https://github.com/harpribot/nlp-metrics
"""
from bleu.bleu import Bleu
from meteor.meteor import Meteor
from rouge.rouge import Rouge
from cider.cider import Cider

import glob


def find_ngrams(input_list, n):
    return zip(*[input_list[i:] for i in range(n)])


def compute_distinct_n_gram(gen_response_words_corpus, n):
    ngrams = find_ngrams(gen_response_words_corpus, n)
    dist_ngrams = set(ngrams)
    return float(len(dist_ngrams)) / float(len(ngrams))


def load_textfiles(references, hypothesis):
    hypo = {idx: [lines.strip()] for (idx, lines) in enumerate(hypothesis)}
    # take out newlines before creating dictionary
    raw_refs = [map(str.strip, r) for r in zip(references)]
    refs = {idx: rr for idx, rr in enumerate(raw_refs)}
    # sanity check that we have the same number of references as hypothesis
    if len(hypo) != len(refs):
        raise ValueError("There is a sentence number mismatch between the inputs")
    return refs, hypo


def score(ref, hypo):
    """
    ref, dictionary of reference sentences (id, sentence)
    hypo, dictionary of hypothesis sentences (id, sentence)
    score, dictionary of scores
    """
    #print('setting up scorers...')
    scorers = [
        (Bleu(4), ["Bleu_1", "Bleu_2", "Bleu_3", "Bleu_4"]),
        #(Meteor(), "METEOR"), # hidde currently due to slow speed
        (Rouge(), "ROUGE_L")
        #(Cider(), "CIDEr")
    ]
    final_scores = {}
    for scorer, method in scorers:
        score, scores = scorer.compute_score(ref, hypo)
        if type(score) == list:
            for m, s in zip(method, score):
                final_scores[m] = s
        else:
            final_scores[method] = score
    # print('final_scores: ', final_scores)
    return final_scores

def score_given_a_sequence_pair(ground_truth_seq, model_generate_seq):
    '''
    :param ground_truth_seq: the ground truth sequence
    :param model_generate_seq: the sequence generated by the model
    :return: a score_dict containing the corresponding BLUE-1, BLUE-2, BLUE-3, BLUE-4 and ROUGE_L scores
    '''
    reference = {'id1': [ground_truth_seq]}
    hypothesis = {'id1': [model_generate_seq]}
    score_dict = score(reference, hypothesis)
    return score_dict

def score_given_multi_sequence_pairs(ground_truth_seqs, model_generate_seqs):
    '''
    :param ground_truth_seqs: dictionary of reference sentences (id, sentence)
    :param model_generate_seqs: dictionary of hypothesis sentences (id, sentence)
    :return: dictionary of scores
    '''
    score_dict = score(ground_truth_seqs, model_generate_seqs)
    return score_dict

if __name__ == '__main__':
    ground_truth_seq = 'Definitely it is very satisfying'
    model_generate_seq = 'NLP Project is challenging but it is very satisfying'
    print ('ground truth seq: ', ground_truth_seq)
    print ('model generate seq: ', model_generate_seq)

    score_dict = score_given_a_sequence_pair(ground_truth_seq, model_generate_seq)

    print ('score_dict: ', score_dict)

    import nltk

    hypothesis = nltk.word_tokenize(model_generate_seq)
    reference = nltk.word_tokenize(ground_truth_seq)
    # there may be several references
    BLEUscore_Sent = nltk.translate.bleu_score.sentence_bleu([reference], hypothesis)
    print ('BLUEscore from NLTK (Sent_BLUE): ', BLEUscore_Sent)

    print ('\nNow test cases with 3 sentence pairs ... ')
    # Feed in the directory where the hypothesis summary and true summary is stored
    hyp_file = glob.glob('hypothesis/*')
    ref_file = glob.glob('reference/*')

    BLEU_1 = 0.
    BLEU_2 = 0.
    BLEU_3 = 0.
    BLEU_4 = 0.
    ROUGE_L = 0.
    num_files = 0
    for reference_file, hypothesis_file in zip(ref_file, hyp_file):
        num_files += 1
        print (reference_file, hypothesis_file)

        with open(reference_file) as rf:
            reference = rf.readlines()

        with open(hypothesis_file) as hf:
            hypothesis = hf.readlines()

        # print 'reference', reference
        # print 'hypothesis', hypothesis

        ref, hypo = load_textfiles(reference, hypothesis)

        # print 'ref', ref
        # print 'hypo', hypo
        score_map = score(ref, hypo)
        BLEU_1 += score_map['Bleu_1']
        BLEU_2 += score_map['Bleu_2']
        BLEU_3 += score_map['Bleu_3']
        BLEU_4 += score_map['Bleu_4']
        ROUGE_L += score_map['ROUGE_L']

    # print 'Average Metric Score for All Review Summary Pairs:'
    # print 'Bleu - 1gram:', BLEU_1 / num_files
    # print 'Bleu - 2gram:', BLEU_2 / num_files
    # print 'Bleu - 3gram:', BLEU_3 / num_files
    # print 'Bleu - 4gram:', BLEU_4 / num_files
    # print 'Rouge:', ROUGE_L / num_files

    # toy data for testing
    # references
    # Harshal is a bad boy .
    # Definitely it is very satisfying .
    # Mary did not slap the green witch .

    # hypotheses
    # Harshal is a good boy .
    # NLP Project is hard but very satisfying .
    # NMary no slap witch green .

    list_of_references = [[nltk.word_tokenize('Harshal is a good'
                                              ' boy .')], \
                          [nltk.word_tokenize('Definitely it is very satisfying .')], \
                          [nltk.word_tokenize('Mary did not slap the green witch .')]]
    hypotheses = [nltk.word_tokenize('Harshal is a good boy .'), \
                  nltk.word_tokenize('NLP Project is hard but very satisfying .'), \
                  nltk.word_tokenize('NMary no slap witch green .'),]

    BLEUscore_Corpus = nltk.translate.bleu_score.corpus_bleu(list_of_references, hypotheses)
    print ('BLUEscore from NLTK (Corpus_BLUE): ', BLEUscore_Corpus)