#!/usr/bin/env python

from __future__ import print_function, division

import argparse, os, re, time
import pdb

from gehrmann_rouge_opennmt.rouge_baselines.g_rouge import rouge
from gehrmann_rouge_opennmt.rouge_baselines.util import has_repeat, n_grams
from functools import reduce
import numpy as np


def split_sentences(article, sentence_start_tag='<t>', sentence_end_tag='</t>'):
    bare_sents = re.findall(r'%s (.+?) %s' % (sentence_start_tag, sentence_end_tag), article)
    return bare_sents


# convenient decorator
def register_to_registry(registry):
    def _register(func):
        registry[func.__name__] = func
        return func

    return _register


baseline_registry = {}
register = register_to_registry(baseline_registry)


# baseline methods
@register
def first_sentence(article, sentence_start_tag='<t>', sentence_end_tag='</t>'):
    ''' use sentence tags to output the first sentence of an article as its summary. '''
    sents = split_sentences(article, sentence_start_tag, sentence_end_tag)
    return sents[:1]


@register
def first_three_sentences(article, sentence_start_tag='<t>', sentence_end_tag='</t>'):
    sents = split_sentences(article, sentence_start_tag, sentence_end_tag)
    return sents[:3]


@register
def first_two_sentences(article, sentence_start_tag='<t>', sentence_end_tag='</t>'):
    sents = split_sentences(article, sentence_start_tag, sentence_end_tag)
    return sents[:2]


@register
def verbatim(article, sentence_start_tag='<t>', sentence_end_tag='</t>'):
    sents = split_sentences(article, sentence_start_tag, sentence_end_tag)
    return sents


@register
def pre_sent_tag_verbatim(article):
    sents = article.split('<t>')
    good_sents = []
    for sent in sents:
        sent = sent.strip()
        if len(sent.split()) > 0:
            good_sents.append(sent)
    # print(good_sents)
    return good_sents


@register
def sent_tag_verbatim(article):
    sents = split_sentences(article, '<t>', '</t>')
    # print(sents)
    return sents


@register
def sent_tag_p_verbatim(article):
    bare_article = article.strip()
    bare_article += ' </t>'
    sents = split_sentences(bare_article, '<t>', '</t>')
    # print(sents)
    return sents


@register
def adhoc_old0(article):
    sents = split_sentences(article, '<t>', '</t>')
    good_sents = []
    for sent in sents:
        # Remove <unk>
        tokens = [x for x in sent.split() if x != '<unk>']
        # Ignore length 1 sententces
        if len(tokens) > 1:
            good_sents.append(' '.join(tokens))
    return good_sents


@register
def full(article):
    return [article]


@register
def adhoc_base(article):
    article += ' </t> </t>'
    first_end = article.index(' </t> </t>')
    article = article[:first_end] + ' </t>'
    sents = split_sentences(article)
    good_sents = []
    for sent in sents:
        # Remove <unk>
        tokens = [x for x in sent.split() if x != '<unk>']
        # Ignore length 1 sententces
        if len(tokens) > 1:
            good_sents.append(' '.join(tokens))
    return good_sents


@register
def no_sent_tag(article):
    article = article.strip()
    try:
        if article[-1] != '.':
            article += ' .'
    except:
        article += ' .'
    good_sents = list(re.findall(r'.+?\.', article))
    return good_sents


@register
def second_sentence(article, sentence_start_tag='<t>', sentence_end_tag='</t>'):
    sents = split_sentences(article, sentence_start_tag, sentence_end_tag)
    return sents[1:2]


def baseline_main(args, return_pyrouge_scores=False):
    # Check the presence of target file
    if args.run_rouge or args.run_google_rouge:
        assert args.target is not None, 'Need the path to target file `--target` for ROUGE evaluations.'

    process = baseline_registry[args.method]

    # Read and preprocess generated summary
    n_source = 0
    references = []
    summaries = []
    with open(args.source, 'r') as f:
        for i, article in enumerate(f):
            summary = process(article)
            summaries.append(summary)
            n_source += 1

    mean_num_sent_per_summ = np.mean([len(summ) for summ in summaries])
    assert mean_num_sent_per_summ > 0, "Expect to read > 0 sentences per summary!"

    # Read and preprocess a single candidate reference summary for each example
    if args.run_rouge or args.run_google_rouge:
        n_target = 0
        with open(args.target, 'r') as f:
            for i, article in enumerate(f):
                # For us, method is 'sent_tag_verbatim
                if args.ref_sep:  # pgour added this to handle multiple reference texts
                    # pdb.set_trace()
                    raw_candidates_l = article.split(args.ref_sep)
                    candidates_l = []
                    for raw_candidate in raw_candidates_l:
                        if args.method == "full":
                            candidate = [raw_candidate]
                        else:
                            candidate = split_sentences(raw_candidate)
                        candidates_l.append(candidate)
                    assert len(candidates_l) == args.num_ref, f"len(candidates_l) {len(candidates_l)} mismatches "\
                                                              f"args.num_ref {args.num_ref}"
                    references.append(candidates_l)
                    n_target += 1
                else:
                    if args.method == "full":
                        candidate = [article]
                    else:
                        candidate = split_sentences(article)
                    references.append([candidate])
                    n_target += 1
        # pdb.set_trace()
        mean_num_sent_per_ref = np.mean([len(candidate[0]) for candidate in references])
        assert mean_num_sent_per_ref > 0, "Expect to read > 0 sentences per reference summary!"

        # logger.info(f"read {mean_num_sent_per_summ:.2f} and {mean_num_sent_per_ref:.2f} sentences on average per "
        #             f"generated and system summary.")

        assert n_source == n_target, 'Source and target must have the same number of samples.'

    # Run official ROUGE evaluation
    if args.run_rouge:
        # logger.info("getting rouge")
        from gehrmann_rouge_opennmt.rouge_baselines.util import evaluate_rouge

        # TODO: what is going on here? Why the double assignment?
        rouge_args = rouge_args = [
            '-c', 95,  # 95% confidence intervals, necessary for the dictionary conversion routine
            '-n', 2,  # up to bigram
            '-a',
            '-r', args.n_bootstrap,  # the number of bootstrap samples for confidence bounds
        ]

        if args.stemming:
            # add the stemming flag
            rouge_args += ['-m']

        if args.get_each_score:
            # add the 'per-evaluation scores' flag
            rouge_args += ['-d']

        # evaluate with official ROUGE script v1.5.5
        scores = evaluate_rouge(summaries, references, remove_temp=args.delete, rouge_args=rouge_args,
                                get_each_score=args.get_each_score, temp_dir=args.temp_dir)

        if return_pyrouge_scores:
            # We always return from here, below this line is not important
            return scores

    # Run Google's ROUGE evaluation. Not used by us.
    if args.run_google_rouge:
        # Based on https://github.com/google/seq2seq, modified to support multi-sentence summaries
        t0 = time.time()
        g_scores = rouge(summaries, [candidates[0] for candidates in references])
        dt = time.time() - t0

        g_headers = ['rouge_1/r_score', 'rouge_1/p_score', 'rouge_1/f_score', 'rouge_2/r_score', 'rouge_2/p_score',
                     'rouge_2/f_score', 'rouge_l/r_score', 'rouge_l/p_score', 'rouge_l/f_score']

        print('* evaluated {} samples, took {:.3f}s, averaging {:.3f}s/sample'.format(n_target, dt, dt / n_target))

    # Evaluate self-repetitions
    if args.check_repeats:
        t0 = time.time()
        # Counts
        n_sent_repeats = 0
        ngram_repeats = {2: 0, 4: 0, 8: 0, 16: 0, 32: 0}
        for summary in summaries:
            # Sentence-level repeats
            # Count of samples containing self-repetitions of a full sentence
            n_sent_repeats += has_repeat(summary)

            # N-gram repeats
            for n in ngram_repeats.keys():
                # Respect sentence boundary
                grams = reduce(lambda x, y: x + y, [n_grams(sent.split(), n) for sent in summary], [])
                ngram_repeats[n] += has_repeat(grams)

        dt = time.time() - t0

        print('* portion of samples that contains self-repetitions')

        # Sort the statistics by importance
        str_keys = ['full-sent'] + list(map(lambda n: '%d-gram' % n, sorted(ngram_repeats.keys(), reverse=True)))
        print(','.join(str_keys))
        print("{:.2f}%".format(n_sent_repeats / n_source * 100), end=',\t')
        for n in sorted(ngram_repeats.keys(), reverse=True):
            print("{:.2f}%".format(ngram_repeats[n] / n_source * 100), end=',\t')
        print()

        print('* evaluated {} samples, took {:.3f}s, averaging {:.3f}s/sample'.format(n_source, dt, dt / n_source))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-s', '--source', required=True,
                        help='Path to the tokenized source file. One sample per line with sentence tags.')
    parser.add_argument('-t', '--target', required=False,
                        help='Path to the tokenized target file. One sample per line with sentence tags.')
    parser.add_argument('-m', '--method', default='first_sentence', choices=baseline_registry.keys(),
                        help='Baseline method to use.')
    parser.add_argument('-d', '--delete', action='store_true',
                        help='Delete the temporary files created during evaluation.')
    parser.add_argument('-g', '--google', dest='run_google_rouge', action='store_true',
                        help='Evaluate with the ROUGE implementation from google/seq2seq.')
    parser.add_argument('--no-rouge', dest='run_rouge', action='store_false', help='Skip ROUGE evaluation.')
    parser.add_argument('-r', '--check-repeats', action='store_true', help='Evaluate self repeats.')
    parser.add_argument('--ref_sep', type=str, default=None, help='if there are multiple references per '
                                                    'line in ref file, they are separated by this separator.') # pgour added
    parser.add_argument('--num_ref', type=int, default=1, help='number of ref summaries for each doc (per line in file)')
    # ROUGE arguments
    parser.add_argument('--no-stemming', dest='stemming', action='store_false', help='Turn off stemming in ROUGE.')
    parser.add_argument('--n-bootstrap', type=int, default=1000, help='The number of bootstrap samples used in ROUGE.')
    parser.add_argument('--get_each_score', action='store_true', help='produce separate score of each document-summary')

    args = parser.parse_args()

    # pgour: sanity check
    if args.num_ref != 1:
        assert (args.ref_sep is not None), "if more than 1 ref per summary, expected a --ref_sep"

    baseline_main(args)
