# coding=utf-8
# Copyright 2023 The TensorFlow Datasets Authors.
#
# 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.

"""ParaCrawl (Bitextor) parallel open-source machine translation benchmark."""

import collections

from etils import epath
from tensorflow_datasets.core import utils
import tensorflow_datasets.public_api as tfds

_BENCHMARK_URL = "https://paracrawl.eu/releases.html"

_BASE_DATA_URL_FORMAT_STR = (
    "https://s3.amazonaws.com/web-language-models/"
    "paracrawl/release4/en-{target_lang}.bicleaner07."
    "txt.gz"
)


@utils.memoize()
def _target_languages():
  """Create the sorted dictionary of language codes, and language names.

  Returns:
    The sorted dictionary as an instance of `collections.OrderedDict`.
  """
  langs = {
      "bg": "Bulgarian",
      "cs": "Czech",
      "da": "Danish",
      "de": "German",
      "el": "Greek",
      "es": "Spanish",
      "et": "Estonian",
      "fi": "Finnish",
      "fr": "French",
      "ga": "Irish",
      "hr": "Croatian",
      "hu": "Hungarian",
      "it": "Italian",
      "lt": "Lithuanian",
      "lv": "Latvian",
      "mt": "Maltese",
      "nl": "Dutch",
      "pl": "Polish",
      "pt": "Portuguese",
      "ro": "Romanian",
      "sk": "Slovak",
      "sl": "Slovenian",
      "sv": "Swedish",
  }
  return collections.OrderedDict(sorted(langs.items()))


class ParaCrawlConfig(tfds.core.BuilderConfig):
  """BuilderConfig for ParaCrawl."""

  def __init__(self, *, target_language=None, **kwargs):
    """BuilderConfig for ParaCrawl.

    Args:
      target_language: Target language that will be used to translate to from
        English which is always the source language. It has to contain 2-letter
        coded strings. For example: "se", "hu".
      **kwargs: Keyword arguments forwarded to super.
    """
    # Validate the target language.
    if target_language not in _target_languages():
      raise ValueError("Invalid target language: %s " % target_language)

    # Initialize the base class.
    name = f"en{target_language}"
    description = "Translation dataset from English to %s." % (target_language)
    super(ParaCrawlConfig, self).__init__(
        name=name,
        description=description,
        version=tfds.core.Version("1.2.0"),
        **kwargs,
    )

    # Store the attributes.
    self.target_language = target_language
    self.data_url = _BASE_DATA_URL_FORMAT_STR.format(
        target_lang=target_language
    )


class Builder(tfds.core.GeneratorBasedBuilder):
  """ParaCrawl machine translation dataset."""

  # Version history:
  # 1.0.0: S3 (new shuffling, sharding and slicing mechanism).
  # 0.1.0: Initial version.
  BUILDER_CONFIGS = [
      # The version below does not refer to the version of the released
      # database. It only indicates the version of the TFDS integration.
      ParaCrawlConfig(  # pylint: disable=g-complex-comprehension
          target_language=target_language,
      )
      for target_language in _target_languages()
  ]

  def _info(self):
    target_language = self.builder_config.target_language
    return self.dataset_info_from_configs(
        features=tfds.features.Translation(
            languages=("en", target_language),
        ),
        supervised_keys=("en", target_language),
        homepage=_BENCHMARK_URL,
    )

  def _split_generators(self, dl_manager):
    # Download the data file.
    data_file = dl_manager.download_and_extract(
        {"data_file": self.builder_config.data_url}
    )

    # Return the single split of the data.
    return [
        tfds.core.SplitGenerator(name=tfds.Split.TRAIN, gen_kwargs=data_file)
    ]

  def _generate_examples(self, data_file):
    """This function returns the examples in the raw (text) form."""
    target_language = self.builder_config.target_language

    with epath.Path(data_file).open() as f:
      for idx, line in enumerate(f):
        line_parts = line.strip().split("\t")
        if len(line_parts) != 2:
          msg = (
              "Wrong data format in line {}. The line '{}' does "
              "not have exactly one delimiter."
          ).format(idx, line)
          raise ValueError(msg)
        source, target = line_parts[0].strip(), line_parts[1].strip()
        yield idx, {"en": source, target_language: target}
