# MIT License

# Copyright (c) 2021 Alex Fabbri, Wojciech Kryściński, Bryan
# McCann, Caiming Xiong, Richard Socher, and Dragomir Radev and The HuggingFace
# Team

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# pylint: disable=C0103,W0221,W0106
# Replace summ_eval.data_stats_metric
import logging
from collections import Counter
from multiprocessing import Pool
from typing import Literal

from lighteval.metrics.imports.data_stats_utils import Fragments
from lighteval.utils.imports import Extra, requires


logger = logging.getLogger(__name__)

LANGUAGE_TO_SPACY_MODEL_MAP = {
    "en": "en_core_web_sm",
    "de": "de_core_news_sm",
    "fr": "fr_core_news_sm",
    "it": "it_core_news_sm",
}


class Metric:
    def evaluate_example(self, summary, reference):
        raise NotImplementedError

    def evaluate_batch(self, summaries, references, aggregate=True):
        raise NotImplementedError


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


@requires(Extra.MULTILINGUAL)
class DataStatsMetric(Metric):
    def __init__(
        self,
        n_gram: int = 3,
        n_workers: int = 24,
        case: bool = False,
        tokenize: bool = True,
        language: Literal["en", "de", "fr", "it"] = "en",
    ):
        """
        Data Statistics metric
        Makes use of Newsroom code: \
            https://github.com/lil-lab/newsroom/blob/master/newsroom/analyze/fragments.py
        Calculates extractive statistics such as coverage, density, compression as
            defined in Newsroom paper as well as the percentage of novel n-grams in the
            summary vs the input text and the percentage of n-grams in the summary which are
            repeated

        NOTE: these statistics are meant to be calculated with respect to the source text
            (e.g. news article) as opposed to the reference.

        Args:
            n_gram (int): compute statistics for n-grams up to and including this length.
            n_workers (int): number of processes to use if using multiprocessing.
            case (bool): whether to lowercase input before calculating statistics.
            tokenize (bool): whether to tokenize the input; otherwise assumes that the input
                is a string of space-separated tokens.
            language (Literal["en", "de", "fr", "it"]): the language of the input text. This
                determines the spaCy model used for tokenization. Currently supports English,
                German, French, and Italian.
        """
        import spacy

        self.n_gram = n_gram
        self.n_workers = n_workers
        self.case = case
        self.tokenize = tokenize
        self.language = language
        self.nlp = None

        spacy_model = LANGUAGE_TO_SPACY_MODEL_MAP.get(self.language, "en_core_web_sm")
        try:
            self.nlp = spacy.load(spacy_model)
        except OSError:
            logger.info(f"Downloading the spacy {spacy_model} model\n(don't worry, this will only happen once)")
            from spacy.cli import download

            download(spacy_model)
            self.nlp = spacy.load(spacy_model)

    def evaluate_example(self, summary, input_text):
        if self.tokenize:
            input_text = self.nlp(input_text, disable=["tagger", "parser", "ner", "textcat"])
            input_text = [tok.text for tok in input_text]
            summary = self.nlp(summary, disable=["tagger", "parser", "ner", "textcat"])
            summary = [tok.text for tok in summary]
        fragments = Fragments(summary, input_text, case=self.case)
        coverage = fragments.coverage()
        density = fragments.density()
        compression = fragments.compression()
        score_dict = {"coverage": coverage, "density": density, "compression": compression}
        tokenized_summary = fragments.summary
        tokenized_text = fragments.text
        score_dict["summary_length"] = len(tokenized_summary)
        for i in range(1, self.n_gram + 1):
            input_ngrams = list(find_ngrams(tokenized_text, i))
            summ_ngrams = list(find_ngrams(tokenized_summary, i))
            input_ngrams_set = set(input_ngrams)
            summ_ngrams_set = set(summ_ngrams)
            intersect = summ_ngrams_set.intersection(input_ngrams_set)
            try:
                score_dict[f"percentage_novel_{i}-gram"] = (len(summ_ngrams_set) - len(intersect)) / float(
                    len(summ_ngrams_set)
                )
                ngramCounter = Counter()
                ngramCounter.update(summ_ngrams)
                repeated = [key for key, val in ngramCounter.items() if val > 1]
                score_dict[f"percentage_repeated_{i}-gram_in_summ"] = len(repeated) / float(len(summ_ngrams_set))
            except ZeroDivisionError:
                continue
        return score_dict

    def evaluate_batch(self, summaries, input_texts, aggregate=True):
        corpus_score_dict = Counter()
        p = Pool(processes=self.n_workers)
        results = p.starmap(self.evaluate_example, zip(summaries, input_texts))
        p.close()
        if aggregate:
            [corpus_score_dict.update(x) for x in results]
            for key in corpus_score_dict.keys():
                corpus_score_dict[key] /= float(len(input_texts))
            return corpus_score_dict
        else:
            return results

    @property
    def supports_multi_ref(self):
        return False
