# coding=utf-8
# Copyright 2020 The TensorFlow Datasets Authors and the HuggingFace 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.

# Lint as: python3
"""The FEVER NLI Corpus."""

from __future__ import absolute_import, division, print_function

import json
import os

import datasets

_CITATION = """\
@inproceedings{thorne-etal-2018-fever,
    title = "{FEVER}: a Large-scale Dataset for Fact Extraction and {VER}ification",
    author = "Thorne, James  and
      Vlachos, Andreas  and
      Christodoulopoulos, Christos  and
      Mittal, Arpit",
    booktitle = "Proceedings of the 2018 Conference of the North 
    {A}merican Chapter of the Association for Computational Linguistics:
    Human Language Technologies, Volume 1 (Long Papers)",
    month = jun,
    year = "2018",
    address = "New Orleans, Louisiana",
    publisher = "Association for Computational Linguistics",
    url = "https://www.aclweb.org/anthology/N18-1074",
    doi = "10.18653/v1/N18-1074",
    pages = "809--819",
    abstract = "In this paper we introduce a new publicly available dataset for verification against textual sources, 
    FEVER: Fact Extraction and VERification. 
    It consists of 185,445 claims generated by altering sentences extracted from Wikipedia 
    and subsequently verified without knowledge of the sentence they were derived from.
    The claims are classified as Supported, Refuted or NotEnoughInfo by annotators achieving 0.6841 in Fleiss kappa.
    For the first two classes, the annotators also recorded the sentence(s) forming the necessary evidence for 
    their judgment. To characterize the challenge of the dataset presented, we develop a pipeline approach and 
    compare it to suitably designed oracles. The best accuracy we achieve on labeling a claim accompanied by the 
    correct evidence is 31.87{\\%}, while if we ignore the evidence we achieve 50.91{\\%}. 
    Thus we believe that FEVER is a challenging testbed that will help stimulate progress on 
    claim verification against textual sources.",
}
"""

_DESCRIPTION = """\
Each data point in the original FEVER dataset is a textual claim paired with a label 
(support / refute / not enough information) depending on whether the claim can be verified by the Wikipedia.
For examples with support and refute labels in the training set and dev set, ground truth location of the evidence 
in the Wikipedia is also provided. (Please refer to [the original paper](https://arxiv.org/abs/1803.05355) for details)
"""


class FeverConfig(datasets.BuilderConfig):
    """BuilderConfig for Fever."""

    def __init__(self, **kwargs):
        """BuilderConfig for Fever.
            Args:
        .
              **kwargs: keyword arguments forwarded to super.
        """
        super(FeverConfig, self).__init__(version=datasets.Version("0.2.0", ""), **kwargs)


class Fever(datasets.GeneratorBasedBuilder):
    # noinspection SpellCheckingInspection
    """Fever: Fact Extraction and VERification. Version 0.2."""

    BUILDER_CONFIGS = [
        FeverConfig(
            name="plain_text",
            description="Plain text",
        ),
    ]

    def _info(self):
        return datasets.DatasetInfo(
            description=_DESCRIPTION,
            features=datasets.Features(
                {
                    "id": datasets.Value("int64"),
                    "claim": datasets.Value("string"),
                    "evidence": datasets.Value("string"),
                    "label": datasets.features.ClassLabel(names=["SUPPORTS", "NOT ENOUGH INFO", "REFUTES"]),
                }
            ),
            # No default supervised_keys (as we have to pass both premise
            # and hypothesis as input).
            supervised_keys=None,
            homepage="https://fever.ai/",
            citation=_CITATION,
        )

    def _vocab_text_gen(self, filepath):
        for _, ex in self._generate_examples(filepath):
            yield " ".join([ex["evidence"], ex["claim"]])

    def _split_generators(self, dl_manager):
        train_path, validation_path = dl_manager.download([
            "https://www.dropbox.com/s/v1a0depfg7jp90f/fever.train.jsonl?dl=1",
            "https://www.dropbox.com/s/bdwf46sa2gcuf6j/fever.dev.jsonl?dl=1"
        ])
        # train_path = os.path.join(train_path, "train_fitems.jsonl")
        # validation_path = os.path.join(validation_path, "dev_fitems.jsonl")

        # Since the FEVER NLI dataset doesn't have labels for the dev set, we also download the original
        # FEVER dev set and match example CIDs to obtain labels.
        # orig_dev_path = dl_manager.download(
        #     "https://s3-eu-west-1.amazonaws.com/fever.public/shared_task_dev.jsonl"
        # )
        # id_to_label = {}
        # with open(orig_dev_path, 'rb') as f:
        #     for idx, line in enumerate(f):
        #         line = line.strip().decode('utf-8')
        #         json_obj = json.loads(line)
        #         if "id" not in json_obj:
        #             print("FEVER dev dataset is missing ID.")
        #             continue
        #         if "label" not in json_obj:
        #             print("FEVER dev dataset is missing label.")
        #             continue
        #         id_to_label[int(json_obj["id"])] = json_obj["label"]
        # self.id_to_label = id_to_label
        # print(len(self.id_to_label))

        return [
            datasets.SplitGenerator(name=datasets.Split.TRAIN, gen_kwargs={"filepath": train_path}),
            datasets.SplitGenerator(name="validation", gen_kwargs={"filepath": validation_path}),
            # datasets.SplitGenerator(name="test", gen_kwargs={"filepath": test_path}),
        ]

    def _generate_examples(self, filepath):
        """Generate mnli examples.
        Args:
          filepath: a string
        Yields:
          dictionaries containing "claim", "evidence" and "label" strings
        """
        with open(filepath, 'rb') as f:
            for idx, line in enumerate(f):
                line = line.strip().decode('utf-8')
                json_obj = json.loads(line)
                # if json_obj["label"] == "hidden":
                #     key = int(json_obj["cid"])
                #     if key not in self.id_to_label:
                #         continue
                #     json_obj["label"] = self.id_to_label[key]

                # Works for both splits even though dev has some extra human labels.
                yield idx, {
                    "id": idx,  # Some IDs show up as duplicates, but we need them to be unique
                    "claim": json_obj["claim"],
                    "evidence": json_obj["evidence"],
                    "label": json_obj["gold_label"]
                }
