# coding=utf-8
# Copyright 2022 DataLab 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.
import datalabs
from datalabs import get_task, TaskType

_CITATION = """\
@article{liu2017adversarial,
title={Adversarial Multi-task Learning for Text Classification},
author={Liu, Pengfei and Qiu, Xipeng and Huang, Xuanjing},
journal={arXiv preprint arXiv:1704.05742},
year={2017}
}
"""

# You can copy an official description
_DESCRIPTION = """\
This dataset is used in the paper of adversarial multi-task learning in text classification including 16 different fields.
"""

_HOMEPAGE = "http://pfliu.com/paper/adv-mtl.html"
# TODO: Add the licence for the dataset here if you can find it
_LICENSE = "N/A"
# The HuggingFace dataset library don't host the datalab but only point to the original files
# This can be an arbitrary nested dict/list of URLs (see below in `_split_generators` method)
_URLs = "https://raw.githubusercontent.com/ShiinaHiiragi/multi-task-dataset/master/{}.task.{}"
_LANGUAGES = ["en"]
_TASK_TEMPLATES = [get_task(TaskType.sentiment_classification)()]

fields = [
    "apparel",
    "baby",
    "books",
    "camera_photo",
    "dvd",
    "electronics",
    "health_personal_care",
    "imdb",
    "kitchen_housewares",
    "magazines",
    "MR",
    "music",
    "software",
    "sports_outdoors",
    "toys_games",
    "video",
]


class AdvMtl(datalabs.GeneratorBasedBuilder):
    BUILDER_CONFIGS = [
        datalabs.BuilderConfig(
            name="{}".format(domain), version=datalabs.Version("1.0.0")
        )
        for domain in fields
    ]


    def _info(self):
        features = datalabs.Features(
            {
                "text": datalabs.Value("string"),
                "label": datalabs.features.ClassLabel(names=["positive", "negative"]),
            }
        )
        return datalabs.DatasetInfo(
            # This is the description that will appear on the datalab page.
            description=_DESCRIPTION,
            # This defines the different columns of the dataset and their types
            features=features,  # Here we define them above because they are different between the two configurations
            # If there's a common (input, target) tuple from the features,
            # specify them here. They'll be used if as_supervised=True in
            # builder.as_dataset.
            supervised_keys=None,
            # Homepage of the dataset for documentation
            homepage=_HOMEPAGE,
            # License for the dataset if available
            license=_LICENSE,
            # Citation for the dataset
            citation=_CITATION,
            task_templates=_TASK_TEMPLATES,
            languages=_LANGUAGES,
        )

    def _split_generators(self, dl_manager):
        domain = str(self.config.name)
        return [
            datalabs.SplitGenerator(
                name=datalabs.Split.TRAIN,
                gen_kwargs={
                    "filepath": dl_manager.download_and_extract(_URLs.format(domain, "train")),
                },
            ),
            datalabs.SplitGenerator(
                name=datalabs.Split.TEST,
                gen_kwargs={
                    "filepath": dl_manager.download_and_extract(_URLs.format(domain, "test")),
                },
            ),
        ]


    def _generate_examples(
        self,
        filepath
    ):
        """Yields examples as (key, example) tuples."""
        # This method handles input defined in _split_generators to yield (key, example) tuples from the dataset.
        # The `key` is here for legacy reason (tfds) and is not important in itself.
        with open(filepath, "rb") as f:
            for id_, row in enumerate(f):
                row = row.decode("utf-8", "ignore")
                datas = row.split("\t")
                text = datas[1]
                label = "positive" if datas[0] == "1" else "negative"
                yield id_, {
                    "text": text,
                    "label": label,
                }
