# 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 datasets
import evaluate

# TODO: Add BibTeX citation
from ocr_evaluation.evaliate.metrics import evaluate_by_words
from ocr_evaluation.ocr.fiftyone import FiftyOneOcr

_CITATION = """\
@InProceedings{huggingface:module,
title = {Iliauni ICC OCR Evaluation},
authors={},
year={2022}
}
"""

# TODO: Add description of the module here
_DESCRIPTION = """\
Better OCR evaluation metric that enables to evaluate OCR results in various ways. It is robust in a way that
it matches the words using their bounding boxes instead of using plain edit distance matching between two texts.
Elaborate more on this later.
"""

# TODO: Add description of the arguments of the module here
_KWARGS_DESCRIPTION = """
Calculates how good are predictions given some references, using certain scores
Args:
    predictions: list of OCR detections in FiftyOne dataset format.
    references: list of OCR detections in FiftyOne dataset format.
Returns:
    evaluation_results: dictionary containing multiple metrics
Examples:
    Examples should be written in doctest format, and should illustrate how
    to use the function.

    >>> dataset = load_dataset("anz2/iliauni_icc_georgian_ocr", use_auth_token="<auth token here>")
    >>> sample = dataset['test'][0]
    >>> ocr_evaluator = evaluate.load("iliauniiccocrevaluation")
    >>> results = ocr_evaluator.compute(references=[sample], predictions=[0, 1])
    >>> print(results)
    {'accuracy': 1.0}
"""


@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
class IliauniIccOCREvaluation(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(
                {
                    "id": datasets.Value("string"),
                    "filepath": datasets.Value("string"),
                    "tags": datasets.Sequence(datasets.Value("string")),
                    "metadata": datasets.Features(
                        {
                            "size_bytes": datasets.Value("int32"),
                            "mime_type": datasets.Value("string"),
                            "width": datasets.Value("int32"),
                            "height": datasets.Value("int32"),
                            "num_channels": datasets.Value("int32"),
                            "author": datasets.Value("string"),
                            "category": datasets.Value("string"),
                            "document_name": datasets.Value("string"),
                            "source": datasets.Value("string"),
                            "year": datasets.Value("int32")
                        }
                    ),
                    "_media_type": datasets.Value("string"),
                    "_rand": datasets.Value("string"),
                    "detections": datasets.Features(
                        {
                            "detections": datasets.Sequence(
                                datasets.Features(
                                    {
                                        "id": datasets.Value("string"),
                                        "attributes": datasets.Sequence(datasets.Value("string")),
                                        "tags": datasets.Value("string"),
                                        "label": datasets.Value("string"),
                                        "bounding_box": datasets.Sequence(datasets.Value("float32")),
                                        "confidence": datasets.Value("float32"),
                                        "index": datasets.Value("int32"),
                                        "page": datasets.Value("int32"),
                                        "block": datasets.Value("int32"),
                                        "paragraph": datasets.Value("int32"),
                                        "word": datasets.Value("int32"),
                                        "text": datasets.Value("string"),
                                    }
                                )
                            )
                        }
                    ),
                    "image": datasets.Image(),
                }
            ),
            # Homepage of the module for documentation
            homepage="http://module.homepage",
            # Additional links to the codebase or references
            codebase_urls=["http://github.com/path/to/codebase/of/new_module"],
            reference_urls=["http://path.to.reference.url/new_module"]
        )

    def _download_and_prepare(self, dl_manager):
        """Optional: download external resources useful to compute the scores"""
        # TODO: Download external resources if needed
        pass

    def _compute(self, predictions, references):
        """Returns the scores"""

        assert len(predictions) == len(references)
        eval_results = []
        for prediction, reference in zip(predictions, references):
            prediction_df = FiftyOneOcr(data=prediction).get_word_annotations(convert_bbox=True)
            reference_df = FiftyOneOcr(data=reference).get_word_annotations(convert_bbox=True)

            eval_result = evaluate_by_words(prediction_df, reference_df, pref1="Pred_", pref2="Tar_")
            eval_results.append(eval_result)

        return eval_results