# coding=utf-8
# Copyright 2021 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.

"""asset dataset for tuning and evaluating text simplification models."""

import csv
import tensorflow as tf
import tensorflow_datasets.public_api as tfds

_CITATION = """\
@inproceedings{alva-manchego-etal-2020-asset,
    title = "{ASSET}: {A} Dataset for Tuning and Evaluation of Sentence \
Simplification Models with Multiple Rewriting Transformations",
    author = "Alva-Manchego, Fernando  and
      Martin, Louis  and
      Bordes, Antoine  and
      Scarton, Carolina  and
      Sagot, Benoit  and
      Specia, Lucia",
    booktitle = "Proceedings of the 58th Annual Meeting of the Association \
for Computational Linguistics",
    month = jul,
    year = "2020",
    address = "Online",
    publisher = "Association for Computational Linguistics",
    url = "https://www.aclweb.org/anthology/2020.acl-main.424",
    pages = "4668--4679",
}
"""

_DESCRIPTION = """\
ASSET is a dataset for evaluating Sentence Simplification systems with \
multiple rewriting transformations, as described in "ASSET: A Dataset for \
Tuning and Evaluation of Sentence Simplification Models with Multiple \
Rewriting Transformations." The corpus is composed of 2000 validation \
and 359 test original sentences that were each simplified 10 times by \
different annotators. The corpus also contains human judgments of meaning \
preservation, fluency and simplicity for the outputs of several automatic \
text simplification systems.
"""

_HOMEPAGE = 'https://github.com/facebookresearch/asset'

_LICENSE = 'Creative Common Attribution-NonCommercial 4.0 International'

_URL_LIST = [
    ('human_ratings.csv',
     'https://raw.githubusercontent.com/facebookresearch/asset/main/human_ratings/human_ratings.csv'
    ),
    ('asset.valid.orig',
     'https://raw.githubusercontent.com/facebookresearch/asset/main/dataset/asset.valid.orig'
    ),
    ('asset.test.orig',
     'https://raw.githubusercontent.com/facebookresearch/asset/main/dataset/asset.test.orig'
    ),
]

# https://github.com/facebookresearch/asset/raw/master/dataset/

_URL_LIST += [
    (  # pylint:disable=g-complex-comprehension
        f'asset.{spl}.simp.{i}',
        f'https://raw.githubusercontent.com/facebookresearch/asset/main/dataset/asset.{spl}.simp.{i}',
    ) for spl in ['valid', 'test'] for i in range(10)
]

_URLs = dict(_URL_LIST)


class Asset(tfds.core.GeneratorBasedBuilder):
  """DatasetBuilder for asset dataset."""

  VERSION = tfds.core.Version('1.0.0')
  RELEASE_NOTES = {
      '1.0.0': 'Initial release.',
  }

  BUILDER_CONFIGS = [
      tfds.core.BuilderConfig(
          name='simplification',
          description='A set of original sentences aligned with 10 possible simplifications for each.'
      ),
      tfds.core.BuilderConfig(
          name='ratings',
          description='Human ratings of automatically produced text simplification.'
      )
  ]

  DEFAULT_CONFIG_NAME = 'simplification'

  def _info(self) -> tfds.core.DatasetInfo:
    """Returns the dataset metadata."""

    if self.builder_config.name == 'simplification':
      features = tfds.features.FeaturesDict({
          'original': tfds.features.Text(),
          'simplifications': tfds.features.Sequence(tfds.features.Text())
      })
    else:
      features = tfds.features.FeaturesDict({
          'original':
              tfds.features.Text(),
          'simplification':
              tfds.features.Text(),
          'original_sentence_id':
              tf.int32,
          'aspect':
              tfds.features.ClassLabel(
                  names=['meaning', 'fluency', 'simplicity']),
          'worker_id':
              tf.int32,
          'rating':
              tf.int32,
      })

    return tfds.core.DatasetInfo(
        builder=self,
        description=_DESCRIPTION,
        features=features,
        supervised_keys=None,
        homepage=_HOMEPAGE,
        citation=_CITATION,
    )

  def _split_generators(self, dl_manager: tfds.download.DownloadManager):
    """Returns SplitGenerators."""

    data_dir = dl_manager.download_and_extract(_URLs)
    if self.builder_config.name == 'simplification':
      return {
          'validation':
              self._generate_examples(filepaths=data_dir, split='valid'),
          'test':
              self._generate_examples(filepaths=data_dir, split='test'),
      }
    else:
      return {
          'full': self._generate_examples(filepaths=data_dir, split='full'),
      }

  def _generate_examples(self, filepaths, split):
    """Yields examples."""

    if self.builder_config.name == 'simplification':
      files = [tf.io.gfile.GFile(filepaths[f'asset.{split}.orig'])] + [
          tf.io.gfile.GFile(filepaths[f'asset.{split}.simp.{i}'])
          for i in range(10)
      ]
      for id_, lines in enumerate(zip(*files)):
        yield id_, {
            'original': lines[0].strip(),
            'simplifications': [line.strip() for line in lines[1:]]
        }
    else:
      with tf.io.gfile.GFile(filepaths['human_ratings.csv']) as f:
        reader = csv.reader(f, delimiter=',')
        for id_, row in enumerate(reader):
          if id_ == 0:
            keys = row[:]
          else:
            res = {}
            for k, v in zip(keys, row):
              res[k] = v
            for k in ['original_sentence_id', 'worker_id', 'rating']:
              res[k] = int(res[k])
            yield (id_ - 1), res
