"""TODO(coqa): Add a description here."""

from __future__ import absolute_import, division, print_function

import csv
import json
import os

import nlp


# TODO(coqa): BibTeX citation
_CITATION = """\
@InProceedings{SivaAndAl:Coca,
       author = {Siva, Reddy and Danqi, Chen and  Christopher D., Manning},
        title = "{WikiQA: A Challenge Dataset for Open-Domain Question Answering}",
      journal = { arXiv},
         year = 2018,

}
"""

# TODO(coqa):
_DESCRIPTION = """\
CoQA: A Conversational Question Answering Challenge
"""

_TRAIN_DATA_URL = "https://nlp.stanford.edu/data/coqa/coqa-train-v1.0.json"
_DEV_DATA_URL = "https://nlp.stanford.edu/data/coqa/coqa-dev-v1.0.json"

#
# class CoqaConfig(nlp.BuilderConfig):
#   """BuilderConfig for Coqa."""
#
#   def __init__(self,
#                story,
#                source,
#                questions,
#                answers,
#                citation,
#                description,
#                additional_answers=None,
#                **kwargs):
#     """BuilderConfig for Coca.
#
#     Args:
#       story: `text`,  context
#       source: `text`, source of the story
#       questions `Sequence` set of questions
#       answers: `Sequence` set of answers to the questions
#       data_url: `string`, url to download the  file from
#       citation: `string`, citation for the data set
#      additional_answers: `Sequence`, in the dev set questions have also set of additional answers
#       **kwargs: keyword arguments forwarded to super.
#     """
#     super(CoqaConfig, self).__init__(
#         version=nlp.Version(
#             "1.0.0",
#             "New split API (https://tensorflow.org/datasets/splits)"),
#         **kwargs)
#     self.story = story
#     self.source = source
#     self.questions = questions
#     self.answers = answers
#     self.additional_answers = additional_answers
#     self.citation = citation
#     self.description = description


class Coqa(nlp.GeneratorBasedBuilder):
    """TODO(coqa): Short description of my dataset."""

    # TODO(coqa): Set up version.
    VERSION = nlp.Version("1.0.0")
    # BUILDER_CONFIGS = CoqaConfig(
    #     story= 'story',
    #     source='source',
    #     questions='questions',
    #     answers='answers',
    #     additional_answers='additional_answers',
    #     description= _DESCRIPTION,
    #     citation= _CITATION
    #
    # )
    def _info(self):
        # TODO(coqa): Specifies the nlp.DatasetInfo object
        return nlp.DatasetInfo(
            # This is the description that will appear on the datasets page.
            description=_DESCRIPTION,
            # nlp.features.FeatureConnectors
            features=nlp.Features(
                {
                    "source": nlp.Value("string"),
                    "story": nlp.Value("string"),
                    "questions": nlp.features.Sequence({"input_text": nlp.Value("string"),}),
                    "answers": nlp.features.Sequence(
                        {
                            "input_text": nlp.Value("string"),
                            "answer_start": nlp.Value("int32"),
                            "answer_end": nlp.Value("int32"),
                        }
                    ),
                    # ##the foloowing feature allows to take into account additional answers in the validation set
                    # 'additional_answers': nlp.features.Sequence({
                    #         "input_texts": nlp.Value('int32'),
                    #         "answers_start": nlp.Value('int32'),
                    #         "answers_end": nlp.Value('int32')
                    #     }),
                    # These are the features of your dataset like images, labels ...
                }
            ),
            # 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="https://stanfordnlp.github.io/coqa/",
            citation=_CITATION,
        )

    def _split_generators(self, dl_manager):
        """Returns SplitGenerators."""
        # TODO(coqa): Downloads the data and defines the splits
        # dl_manager is a nlp.download.DownloadManager that can be used to
        # download and extract URLs
        urls_to_download = {"train": _TRAIN_DATA_URL, "dev": _DEV_DATA_URL}
        downloaded_files = dl_manager.download_and_extract(urls_to_download)

        return [
            nlp.SplitGenerator(
                name=nlp.Split.TRAIN, gen_kwargs={"filepath": downloaded_files["train"], "split": "train"}
            ),
            nlp.SplitGenerator(
                name=nlp.Split.VALIDATION, gen_kwargs={"filepath": downloaded_files["dev"], "split": "validation"}
            ),
        ]

    def _generate_examples(self, filepath, split):
        """Yields examples."""
        # TODO(coqa): Yields (key, example) tuples from the dataset
        with open(filepath) as f:
            data = json.load(f)
            for row in data["data"]:
                questions = [question["input_text"] for question in row["questions"]]
                story = row["story"]
                source = row["source"]
                answers_start = [answer["span_start"] for answer in row["answers"]]
                answers_end = [answer["span_end"] for answer in row["answers"]]
                answers = [answer["input_text"] for answer in row["answers"]]
                # add_answers = row['additional_answers']
                # add_input_tests = []
                # add_start_answers = []
                # add_end_answers = []
                # for key in add_answers:
                #     add_answers_key = add_answers[key]
                #     add_input_tests.append([add_answer['input_text'] for add_answer in add_answers_key])
                #     add_start_answers.append([add_answer['span_start'] for add_answer in add_answers_key])
                #     add_end_answers.append([add_answer['span_end'] for add_answer in add_answers_key])
                yield row["id"], {
                    "source": source,
                    "story": story,
                    "questions": {"input_text": questions,},
                    "answers": {"input_text": answers, "answer_start": answers_start, "answer_end": answers_end}
                    # 'additional_answers': {
                    #     "input_texts": add_input_tests ,
                    #     "answers_start": add_start_answers,
                    #     "answers_end": add_end_answers,
                    # }
                }
