# 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
import ast
import math


# 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 compute a coordinate accuracy between coordinates generated by an LM and a golden one. A pair of coordinates is considered correct if its haversine distance to the golden coordinates is inferior to a threeshold d.
"""


# TODO: Add description of the arguments of the module here
_KWARGS_DESCRIPTION = """
Calculates how good are predictions given some references, using certain scores
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 two floats corresponding to the latitude and longitude of the ground truth (eg. [12.8, 76.9]).
Returns:
    accuracy: 1 if coordinates predicted are d distant from gold ones, O otherwise.
Examples:
    >>> my_new_module = evaluate.load("rfr2003/coord_eval")
    >>> results = my_new_module.compute(generations=["(12.7, 67.8)", "(16.7, 89.6)"], golds=[[12.7, 67.8], [10.9, 80.6]], d_range=20)
    >>> print(results)
    {'coord_accuracy': 0.5}
"""


@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
class Coord_eval(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"""
        # TODO: Download external resources if needed
        pass

    def _haversine_distance(self, coord1, coord2):
        lat1, lon1 = coord1
        lat2, lon2 = coord2
    
        # Convert degrees to radians
        lat1, lon1, lat2, lon2 = map(math.radians, [lat1, lon1, lat2, lon2])
    
        dlat = lat2 - lat1
        dlon = lon2 - lon1
    
        # Haversine formula
        a = math.sin(dlat/2)**2 + math.cos(lat1) * math.cos(lat2) * math.sin(dlon/2)**2
        c = 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))
        R = 6371.0
    
        return R * c

    def _accuracy_coord(self, gen, gold, d_range=20):
        # 1 if gen is in a range of d_range km of gold
        d = self._haversine_distance(gold, gen)
        return int(d <= d_range)

    def _compute(self, generations, golds, d_range=20):
        assert len(generations) == len(golds)
        assert isinstance(golds, list)
    
        correct, total = 0, 0
    
        for gen, gold in zip(generations, golds):
            # Each gold must be at the format : [lat, long]
            assert len(gold) == 2
            f_gold = (float(gold[0]), float(gold[1]))
    
            try:
                f_ans = ast.literal_eval(gen)
                correct += self._accuracy_coord(f_ans, f_gold, d_range)
                
            except:
                pattern = r'[\(\[]\s*(\d+((,|\s|\.)*\d+)*)\s*[, ]\s*(\d+((,|\s|\.)*\d+)*)\s*[\)\]]'
                matches = re.findall(pattern, gen)
                
                if matches:
                    match = matches[0]
                    f_ans = (float(match[0].replace(',', '.').replace(' ', '')), float(match[3].replace(',', '.').replace(' ', '')))
                    correct += self._accuracy_coord(f_ans, f_gold, d_range)
                            
            total += 1

        metrics = {}
        metrics.update({ 
            'coord_ accuracy': correct/total,
        })
    
        return metrics