# Copyright 2020 The HuggingFace Datasets Authors and the current dataset script contributor.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# huggingface packages
import evaluate
import datasets

# faireval functions
from .FairEvalUtils import *

# packages to manage input formats
import importlib
from typing import List, Optional, Union
from seqeval.metrics.v1 import check_consistent_length
from seqeval.scheme import Entities, Token, auto_detect

_CITATION = """\
@inproceedings{ortmann2022,
    title = {Fine-Grained Error Analysis and Fair Evaluation of Labeled Spans},
    author = {Katrin Ortmann},
    url = {https://aclanthology.org/2022.lrec-1.150},
    year = {2022},
    date = {2022-06-21},
    booktitle = {Proceedings of the Language Resources and Evaluation Conference (LREC)},
    pages = {1400-1407},
    publisher = {European Language Resources Association},
    address = {Marseille, France},
    pubstate = {published},
    type = {inproceedings}
}
"""

_DESCRIPTION = """\
New evaluation method that more accurately reflects true annotation quality by ensuring that every error is counted 
only once - avoiding the penalty to close-to-target annotations happening in traditional evaluation. 
In addition to the traditional categories of true positives (TP), false positives (FP), and false negatives 
(FN), the new method takes into account more fine-grained error types: labeling errors (LE), boundary errors (BE), 
and labeling-boundary errors (LBE). 
"""

_KWARGS_DESCRIPTION = """
Outputs the error count (TP, FP, etc.) and resulting scores (Precision, Recall and F1) from a reference list of 
spans compared against a predicted one. The user can choose to see traditional or fair error counts and scores by 
switching the argument 'mode'.
For the computation of the fair metrics from the error count please refer to: https://aclanthology.org/2022.lrec-1.150.pdf
Args:
    predictions: a list of lists of predicted labels, i.e. estimated targets as returned by a tagger.
    references: list of ground truth reference labels. Predicted sentences must have the same number of tokens as the references.
    mode: 'fair' or 'traditional'. Controls the desired output. 'Traditional' is equivalent to seqeval's metrics. The default value is 'fair'.
    error_format: 'count' or 'proportion'. Controls the desired output for TP, FP, BE, LE, etc. 'count' gives the absolute count per parameter. 'proportion' gives the precentage with respect to the total errors that each parameter represents. Default value is 'count'.
    zero_division: which value to substitute as a metric value when encountering zero division. Should be one of [0,1,"warn"]. "warn" acts as 0, but the warning is raised.
    suffix: True if the IOB tag is a suffix (after type) instead of a prefix (before type), False otherwise. The default value is False, i.e. the IOB tag is a prefix (before type).
    scheme: the target tagging scheme, which can be one of [IOB1, IOB2, IOE1, IOE2, IOBES, BILOU]. The default value is None.
Returns:
    A dictionary with:
        - Overall error parameter count (or ratio) and resulting scores.
        - A nested dictionary per label with its respective error parameter count (or ratio) and resulting scores

    If mode is 'traditional', the error parameters shown are the classical TP, FP and FN. If mode is 'fair', TP remain the same,
    FP and FN are shown as per the fair definition and additional errors BE, LE and LBE are shown.
    
Examples:
    >>> faireval = evaluate.load("hpi-dhc/FairEval")
    >>> pred = [['O', 'O', 'B-MISC', 'I-MISC', 'I-MISC', 'I-MISC', 'O', 'B-PER', 'I-PER', 'O']]
    >>> ref =  [['O', 'O', 'O',      'B-MISC', 'I-MISC', 'I-MISC', 'O', 'B-PER', 'I-PER', 'O']]
    >>> results = faireval.compute(predictions=pred, references=ref, mode='fair', error_format='count)
    >>> print(results)
    {'MISC': {'precision': 0.0, 'recall': 0.0, 'f1': 0.0, 'TP': 0,'FP': 0,'FN': 0,'LE': 0,'BE': 1,'LBE': 0},
    'PER': {'precision': 1.0,'recall': 1.0,'f1': 1.0,'TP': 1,'FP': 0,'FN': 0,'LE': 0,'BE': 0,'LBE': 0},
    'overall_precision': 0.6666666666666666,
    'overall_recall': 0.6666666666666666,
    'overall_f1': 0.6666666666666666,
    'TP': 1,
    'FP': 0,
    'FN': 0,
    'LE': 0,
    'BE': 1,
    'LBE': 0}
    """


@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
class FairEvaluation(evaluate.Metric):

    def _info(self):
        return evaluate.MetricInfo(
            # This is the description that will appear on the modules page.
            module_type="metric",
            description=_DESCRIPTION,
            citation=_CITATION,
            inputs_description=_KWARGS_DESCRIPTION,
            # This defines the format of each prediction and reference
            features=datasets.Features({
                "predictions": datasets.Sequence(datasets.Value("string", id="label"), id="sequence"),
                "references": datasets.Sequence(datasets.Value("string", id="label"), id="sequence"),
            }),
            # Homepage of the module for documentation
            homepage="https://huggingface.co/spaces/illorca/fairevaluation",
            # Additional links to the codebase or references
            codebase_urls=["https://github.com/rubcompling/FairEval#acknowledgement"],
            reference_urls=["https://aclanthology.org/2022.lrec-1.150.pdf"]
        )

    def _compute(
            self,
            predictions,
            references,
            suffix: bool = False,
            scheme: Optional[str] = None,
            mode: Optional[str] = 'fair',
            weights: dict = None,
            error_format: Optional[str] = 'count',
            zero_division: Union[str, int] = "warn",
    ):
        """Returns the error parameter counts and scores"""
        # (1) SEQEVAL INPUT MANAGEMENT
        if scheme is not None:
            try:
                scheme_module = importlib.import_module("seqeval.scheme")
                scheme = getattr(scheme_module, scheme)
            except AttributeError:
                raise ValueError(f"Scheme should be one of [IOB1, IOB2, IOE1, IOE2, IOBES, BILOU], got {scheme}")

        y_true = references
        y_pred = predictions

        check_consistent_length(y_true, y_pred)

        if scheme is None or not issubclass(scheme, Token):
            scheme = auto_detect(y_true, suffix)

        true_spans = Entities(y_true, scheme, suffix).entities
        pred_spans = Entities(y_pred, scheme, suffix).entities

        # (2) TRANSFORM FROM SEQEVAL TO FAIREVAL SPAN FORMAT
        true_spans = seq_to_fair(true_spans)
        pred_spans = seq_to_fair(pred_spans)

        # (3) COUNT ERRORS AND CALCULATE SCORES
        total_errors = compare_spans([], [])
        for i in range(len(true_spans)):
            sentence_errors = compare_spans(true_spans[i], pred_spans[i])
            total_errors = add_dict(total_errors, sentence_errors)

        if weights is None and mode == 'weighted':
            print("The chosen mode is \'weighted\', but no weights are given. Setting weights to:\n")
            weights = {"TP": {"TP": 1},
                 "FP": {"FP": 1},
                 "FN": {"FN": 1},
                 "LE": {"TP": 0, "FP": 0.5, "FN": 0.5},
                 "BE": {"TP": 0.5, "FP": 0.25, "FN": 0.25},
                 "LBE": {"TP": 0, "FP": 0.5, "FN": 0.5}}
            print(weights)

        config = {"labels": "all", "eval_method": [mode], "weights": weights,}
        results = calculate_results(total_errors, config)
        del results['conf']

        # (4) SELECT OUTPUT MODE AND REFORMAT AS SEQEVAL-HUGGINGFACE OUTPUT
        # initialize empty dictionary and count errors
        output = {}
        total_trad_errors = results['overall']['traditional']['FP'] + results['overall']['traditional']['FN']
        total_fair_errors = results['overall']['fair']['FP'] + results['overall']['fair']['FN'] + \
                            results['overall']['fair']['LE'] + results['overall']['fair']['BE'] + \
                            results['overall']['fair']['LBE']

        # assert valid options
        assert mode in ['traditional', 'fair', 'weighted'], 'mode must be \'traditional\', \'fair\' or \'weighted\''
        assert error_format in ['count', 'proportion'], 'error_format must be \'count\' or \'proportion\''

        # append entity-level errors and scores
        if mode == 'traditional':
            for k, v in results['per_label'][mode].items():
                if error_format == 'count':
                    output[k] = {'precision': v['Prec'], 'recall': v['Rec'], 'f1': v['F1'], 'TP': v['TP'],
                                 'FP': v['FP'], 'FN': v['FN']}
                elif error_format == 'proportion':
                    output[k] = {'precision': v['Prec'], 'recall': v['Rec'], 'f1': v['F1'], 'TP': v['TP'],
                                 'FP': v['FP'] / total_trad_errors, 'FN': v['FN'] / total_trad_errors}
        elif mode == 'fair' or mode == 'weighted':
            for k, v in results['per_label'][mode].items():
                if error_format == 'count':
                    output[k] = {'precision': v['Prec'], 'recall': v['Rec'], 'f1': v['F1'], 'TP': v['TP'],
                                 'FP': v['FP'], 'FN': v['FN'], 'LE': v['LE'], 'BE': v['BE'], 'LBE': v['LBE']}
                elif error_format == 'proportion':
                    output[k] = {'precision': v['Prec'], 'recall': v['Rec'], 'f1': v['F1'], 'TP': v['TP'],
                                 'FP': v['FP'] / total_fair_errors, 'FN': v['FN'] / total_fair_errors,
                                 'LE': v['LE'] / total_fair_errors, 'BE': v['BE'] / total_fair_errors,
                                 'LBE': v['LBE'] / total_fair_errors}

        # append overall scores
        output['overall_precision'] = results['overall'][mode]['Prec']
        output['overall_recall'] = results['overall'][mode]['Rec']
        output['overall_f1'] = results['overall'][mode]['F1']

        # append overall error counts
        if mode == 'traditional':
            output['TP'] = results['overall'][mode]['TP']
            output['FP'] = results['overall'][mode]['FP']
            output['FN'] = results['overall'][mode]['FN']
            if error_format == 'proportion':
                output['FP'] = output['FP'] / total_trad_errors
                output['FN'] = output['FN'] / total_trad_errors
        elif mode == 'fair' or 'weighted':
            output['TP'] = results['overall'][mode]['TP']
            output['FP'] = results['overall'][mode]['FP']
            output['FN'] = results['overall'][mode]['FN']
            output['LE'] = results['overall'][mode]['LE']
            output['BE'] = results['overall'][mode]['BE']
            output['LBE'] = results['overall'][mode]['LBE']
            if error_format == 'proportion':
                output['FP'] = output['FP'] / total_fair_errors
                output['FN'] = output['FN'] / total_fair_errors
                output['LE'] = output['LE'] / total_fair_errors
                output['BE'] = output['BE'] / total_fair_errors
                output['LBE'] = output['LBE'] / total_fair_errors

        return output


def seq_to_fair(seq_sentences):
    "Transforms input anotated sentences from seqeval span format to FairEval span format"
    out = []
    for seq_sentence in seq_sentences:
        sentence = []
        for entity in seq_sentence:
            span = str(entity).replace('(', '').replace(')', '').replace(' ', '').split(',')
            span = span[1:]
            span[-1] = int(span[-1]) - 1
            span[1] = int(span[1])
            span.append({i for i in range(span[1], span[2] + 1)})
            sentence.append(span)
        out.append(sentence)
    return out