# 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.
"""TODO: Add a description here."""

import evaluate
import datasets
import re
from statistics import median
import numpy as np
import ast

# TODO: Add BibTeX citation
_CITATION = """\
@InProceedings{huggingface:module,
title = {A great new module},
authors={huggingface, Inc.},
year={2020}
}
"""

# TODO: Add description of the module here
_DESCRIPTION = """\
This metric aims to evaluate regression tasks done by LMs.
"""


# TODO: Add description of the arguments of the module here
_KWARGS_DESCRIPTION = """
Calculates Accuracy and Blue-1 between generations and gold answers in a MCQ context.
Args:
    generations: list of predictions to score. Each predictions
        should be a string generated by a LM model.
    golds: list of reference for each prediction. Each
        reference should be a list of floats.
Returns:
    precision:  ,
    recall:
Examples:
    Here is an exemple on how to use the metric:

    >>> metric = evaluate.load("rfr2003/regression_evaluate")
    >>> results = metric.compute(generations=['[150, 0]'], golds=[183, 177, 146, 85, 70, 78, 55, 17, 0, -1, -1])
    >>> print(results)
    {'precision': 4.0, 'recall': 344.0, 'macro-mean': 174.0, 'median macro-mean': 174.0}
"""


@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
class regression_evaluate(evaluate.Metric):
    """TODO: Short description of my evaluation module."""

    def _info(self):
        # TODO: Specifies the evaluate.EvaluationModuleInfo object
        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({
                'generations': datasets.Value('string'),
                'golds': datasets.Sequence(datasets.Value('float32')),
            }),
        )

    def _download_and_prepare(self, dl_manager):
        """Optional: download external resources useful to compute the scores"""
        pass

    def _calculate_pre_rec(self, gens, golds, d):
        dists = []
    
        for gold in golds:
            g_dist = []
            for gen in gens:
                g_dist.append(d(gold, gen))
            if len(g_dist) == 0:
                g_dist.append(100) #penalty if the model doesnt generate anything
            dists.append(g_dist)
    
        dists = np.array(dists)
        precision = np.min(dists, axis=0).sum()
    
        recall = np.min(dists, axis=1).sum()
    
        return precision, recall

    def _compute(self, generations, golds):
        assert len(generations) == len(golds)
        assert isinstance(golds, list)
    
        precisions, recalls, means_pre_rec = [], [], []
    
        for gen, gold in zip(generations, golds):
            f_gold = list(set([float(g) for g in gold]))
            
            f_ans = re.findall(r'\d+(?:,\d+)*(?:\.\d+)?|\d+(?:\.\d+)?', gen)
                
            f_ans = list(set([float(a.replace(',', '')) for a in f_ans])) #get rid of duples values
    
            precision, recall = self._calculate_pre_rec(f_ans, f_gold, lambda x,y: abs(x-y))
    
            precisions.append(precision)
            recalls.append(recall)
            means_pre_rec.append((precision+recall)/2)
            
    
        macro_prec = np.mean(precisions).item()
        macro_rec = np.mean(recalls).item()
    
        metrics = {}
        metrics.update({ 
            'precision': np.mean(precisions).item(),
            'recall': np.mean(recalls).item(),
            'macro-mean': np.mean(means_pre_rec).item(),
            'median macro-mean': median(means_pre_rec)
        })
    
        return metrics