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

"""TriviaQA: A Reading Comprehension Dataset."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import json
import os
from absl import logging
import tensorflow as tf
import tensorflow_datasets.public_api as tfds


_CITATION = """
@article{2017arXivtriviaqa,
       author = {{Joshi}, Mandar and {Choi}, Eunsol and {Weld},
                 Daniel and {Zettlemoyer}, Luke},
        title = "{triviaqa: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension}",
      journal = {arXiv e-prints},
         year = 2017,
          eid = {arXiv:1705.03551},
        pages = {arXiv:1705.03551},
archivePrefix = {arXiv},
       eprint = {1705.03551},
}
"""
_DOWNLOAD_URL = (
    "http://nlp.cs.washington.edu/triviaqa/data/triviaqa-rc.tar.gz")
_TOP_LEVEL_DIRNAME = "qa"
_EVIDENCE_DIRNAME = "evidence"
_TRAIN_FILE_FORMAT = os.path.join(_TOP_LEVEL_DIRNAME, "*-train.json")
_HELDOUT_FILE_FORMAT = os.path.join(_TOP_LEVEL_DIRNAME,
                                    "*test-without-answers.json")
_WEB_EVIDENCE_DIR = os.path.join(_EVIDENCE_DIRNAME, "web")
_WIKI_EVIDENCE_DIR = os.path.join(_EVIDENCE_DIRNAME, "wikipedia")

_DESCRIPTION = """\
TriviaqQA is a reading comprehension dataset containing over 650K
question-answer-evidence triples. TriviaqQA includes 95K question-answer
pairs authored by trivia enthusiasts and independently gathered evidence
documents, six per question on average, that provide high quality distant
supervision for answering the questions.
"""


def _train_data_filenames(tmp_dir):
  return tf.io.gfile.glob(os.path.join(tmp_dir, _TRAIN_FILE_FORMAT))


def _test_data_filenames(tmp_dir):
  return tf.io.gfile.glob(os.path.join(tmp_dir, _HELDOUT_FILE_FORMAT))


def _web_evidence_dir(tmp_dir):
  return tf.io.gfile.glob(os.path.join(tmp_dir, _WEB_EVIDENCE_DIR))


def _wiki_evidence_dir(tmp_dir):
  return tf.io.gfile.glob(os.path.join(tmp_dir, _WIKI_EVIDENCE_DIR))


class TriviaQA(tfds.core.GeneratorBasedBuilder):
  """TriviaQA is a reading comprehension dataset.

  It containss over 650K question-answer-evidence triples.
  """

  VERSION = tfds.core.Version("0.1.0",
                              experiments={tfds.core.Experiment.S3: False})
  SUPPORTED_VERSIONS = [
      tfds.core.Version("1.0.0")
  ]
  # Version history:
  # 1.0.0: S3 (new shuffling, sharding and slicing mechanism).
  # 0.1.0: Initial version.

  def _info(self):
    return tfds.core.DatasetInfo(
        builder=self,
        description=_DESCRIPTION,
        features=tfds.features.FeaturesDict({
            "question":
                tfds.features.Text(),
            "question_id":
                tfds.features.Text(),
            "question_source":
                tfds.features.Text(),
            "entity_pages":
                tfds.features.Sequence({
                    "doc_source": tfds.features.Text(),
                    "file_name": tfds.features.Text(),
                    "title": tfds.features.Text(),
                    "wiki_context": tfds.features.Text(),
                }),
            "search_results":
                tfds.features.Sequence({
                    "description":
                        tfds.features.Text(),
                    "file_name":
                        tfds.features.Text(),
                    "rank":
                        tf.int32,
                    "title":
                        tfds.features.Text(),
                    "url":
                        tfds.features.Text(),
                    "search_context":
                        tfds.features.Text(),
                }),
            "answer":
                tfds.features.FeaturesDict({
                    "aliases":
                        tfds.features.Sequence(tfds.features.Text()),
                    "normalized_aliases":
                        tfds.features.Sequence(tfds.features.Text()),
                    "matched_wiki_entity_name":
                        tfds.features.Text(),
                    "normalized_matched_wiki_entity_name":
                        tfds.features.Text(),
                    "normalized_value":
                        tfds.features.Text(),
                    "type":
                        tfds.features.Text(),
                    "value":
                        tfds.features.Text(),
                }),
        }),

        supervised_keys=None,
        urls=["http://nlp.cs.washington.edu/triviaqa/"],
        citation=_CITATION,
    )

  def _split_generators(self, dl_manager):
    """Returns SplitGenerators."""
    trivia_path = dl_manager.download_and_extract(_DOWNLOAD_URL)

    train_files = _train_data_filenames(trivia_path)
    test_files = _test_data_filenames(trivia_path)
    web_evidence_dir = _web_evidence_dir(trivia_path)
    wiki_evidence_dir = _wiki_evidence_dir(trivia_path)

    return [
        tfds.core.SplitGenerator(
            name=tfds.Split.TRAIN,
            num_shards=100,
            gen_kwargs={"files": train_files,
                        "web_dir": web_evidence_dir,
                        "wiki_dir": wiki_evidence_dir}),
        tfds.core.SplitGenerator(
            name=tfds.Split.TEST,
            num_shards=10,
            gen_kwargs={"files": test_files,
                        "web_dir": web_evidence_dir,
                        "wiki_dir": wiki_evidence_dir}),
    ]

  def _generate_examples(self, files, web_dir, wiki_dir):
    """This function returns the examples."""
    for filepath in files:
      logging.info("generating examples from = %s", filepath)
      with tf.io.gfile.GFile(filepath) as f:
        triviaqa = json.load(f)
        for article in triviaqa["Data"]:
          if "Answer" in article:
            answer = article["Answer"]
            aliases = [alias.strip() for alias in answer["Aliases"]]
            normalized_aliases = [
                alias.strip() for alias in answer["NormalizedAliases"]
            ]
            matched_wiki_entity_name = answer.get("MatchedWikiEntryName",
                                                  "").strip()
            normalized_matched_wiki_entity_name = answer.get(
                "NormalizedMatchedWikiEntryName", "").strip()
            normalized_value = answer["NormalizedValue"].strip()
            type_ = answer["Type"].strip()
            value = answer["Value"].strip()

            answer_dict = {
                "aliases":
                    aliases,
                "normalized_aliases":
                    normalized_aliases,
                "matched_wiki_entity_name":
                    matched_wiki_entity_name,
                "normalized_matched_wiki_entity_name":
                    normalized_matched_wiki_entity_name,
                "normalized_value":
                    normalized_value,
                "type":
                    type_,
                "value":
                    value,
            }
          else:
            answer_dict = {
                "aliases":
                    [],
                "normalized_aliases":
                    [],
                "matched_wiki_entity_name":
                    "<unk>",
                "normalized_matched_wiki_entity_name":
                    "<unk>",
                "normalized_value":
                    "<unk>",
                "type":
                    "",
                "value":
                    "<unk>",
            }

          if "SearchResults" in article:
            descriptions = [search_result["Description"].strip() for
                            search_result in article["SearchResults"]]
            search_file_names = [
                search_result["Filename"].strip() for
                search_result in article["SearchResults"]
            ]
            ranks = [
                search_result["Rank"] for
                search_result in article["SearchResults"]
            ]
            titles = [
                search_result["Title"].strip() for
                search_result in article["SearchResults"]
            ]
            urls = [
                search_result["Url"].strip() for
                search_result in article["SearchResults"]
            ]
            search_contexts = []
            for file_name in search_file_names:
              try:
                search_file = os.path.join(web_dir[0], file_name)
                with tf.io.gfile.GFile(search_file) as f:
                  text = f.read()
                  search_contexts.append(text)
              except (IOError, tf.errors.NotFoundError):
                logging.info("File does not exist, skipping: %s", file_name)
                search_contexts.append("")
          else:
            descriptions = []
            search_file_names = []
            ranks = []
            titles = []
            urls = []
            search_contexts = []

          question = article["Question"].strip()
          question_id = article["QuestionId"]
          question_source = article["QuestionSource"].strip()

          if article["EntityPages"]:
            doc_sources = [
                entitypage["DocSource"] for entitypage in article["EntityPages"]
            ]
            file_names = [
                entitypage["Filename"] for entitypage in article["EntityPages"]
            ]
            wiki_titles = [
                entitypage["Title"] for entitypage in article["EntityPages"]
            ]
            wiki_contexts = []
            for file_name in file_names:
              try:
                wiki_file = os.path.join(wiki_dir[0], file_name)
                with tf.io.gfile.GFile(wiki_file) as f:
                  text = f.read()
                  wiki_contexts.append(text)
              except (IOError, tf.errors.NotFoundError):
                logging.info("File does not exist, skipping: %s", file_name)
                wiki_contexts.append("")
          else:
            doc_sources = []
            file_names = []
            wiki_titles = []
            wiki_contexts = []

          yield question_id, {
              "entity_pages": {
                  "doc_source": doc_sources,
                  "file_name": file_names,
                  "title": wiki_titles,
                  "wiki_context": wiki_contexts,
              },
              "search_results": {
                  "description": descriptions,
                  "file_name": search_file_names,
                  "rank": ranks,
                  "title": titles,
                  "url": urls,
                  "search_context": search_contexts,
              },
              "question": question,
              "question_id": question_id,
              "question_source": question_source,
              "answer": answer_dict,
          }
