import datasets
from nltk.metrics.distance import edit_distance

import evaluate


_DESCRIPTION = """
The Normalized Edit Distance (NED) is a metric used to quantify the dissimilarity between two sequences, typically strings, by measuring the minimum number of editing operations required to transform one sequence into the other, normalized by the length of the longer sequence. 
The NED ranges from 0 to 1, where 0 indicates identical sequences and 1 indicates completely dissimilar sequences.
It is particularly useful in tasks such as spell checking, speech recognition, and OCR. The normalized edit distance can be calculated using the formula:

NED = (1 - (ED(pred, gt) / max(length(pred), length(gt))))

Where:
gt: ground-truth sequence
pred: predicted sequence
ED: Edit Distance, the minimum number of editing operations (insertions, deletions, substitutions) needed to transform one sequence into the other.
"""

_KWARGS_DESCRIPTION = """
Args:
    predictions (`list` of `str`): Predicted labels.
    references (`list` of `str`): Ground truth labels.
Returns:
    ned (`float` or `int`): NED score. Minimum possible value is 0. Maximum possible value is 1.0.
Examples:
    Example 
        >>> ned_metric = evaluate.load("Bekhouche/NED")
        >>> results = ned_metric.compute(references=['text', 'color', 'deep'], predictions=['texts', 'colours', 'deep'])
        >>> print(results)
        {'ned': 0.8380952380952381}
"""

_CITATION = """
@inproceedings{sun2019icdar,
  title={ICDAR 2019 competition on large-scale street view text with partial labeling-RRC-LSVT},
  author={Sun, Yipeng and Ni, Zihan and Chng, Chee-Kheng and Liu, Yuliang and Luo, Canjie and Ng, Chun Chet and Han, Junyu and Ding, Errui and Liu, Jingtuo and Karatzas, Dimosthenis and others},
  booktitle={2019 International Conference on Document Analysis and Recognition (ICDAR)},
  pages={1557--1562},
  year={2019},
  organization={IEEE}
}
"""


@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
class NED(evaluate.Metric):
    def _info(self):
        return evaluate.MetricInfo(
            description=_DESCRIPTION,
            citation=_CITATION,
            inputs_description=_KWARGS_DESCRIPTION,
            features=datasets.Features(
                {
                    "predictions": datasets.Sequence(datasets.Value("string")),
                    "references": datasets.Sequence(datasets.Value("string")),
                }
                if self.config_name == "multilabel"
                else {
                    "predictions": datasets.Value("string"),
                    "references": datasets.Value("string"),
                }
            ),
            reference_urls=["https://huggingface.co/spaces/Bekhouche/NED"],
        )

    def _compute(self, predictions, references):
        ned = 0.0
        if isinstance(predictions, list) and isinstance(references, list):
            ned = sum([self._compute_ned(prediction, reference) for prediction, reference in zip(predictions, references)])/len(predictions)
        elif isinstance(predictions, str) and isinstance(references, str):
            ned = self._compute_ned(predictions, references)
        else:
            raise ValueError("Predictions and references must be either a list[str] or str")
        return {
            "ned": float(
                ned
            )
        }
    
    @staticmethod
    def _compute_ned(prediction, reference):
        if len(reference) == 0 or len(prediction) == 0:
            return 0
        ed = edit_distance(prediction, reference)
        ned = 1 - (ed / max(len(prediction), len(reference)))
        return ned