# Copyright 2017 Google Inc. All Rights Reserved.
#
# 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.
# ==============================================================================
"""Input readers and document/token generators for datasets."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from collections import namedtuple
import csv
import os
import random

# Dependency imports

import tensorflow as tf

from data import data_utils

flags = tf.app.flags
FLAGS = flags.FLAGS

flags.DEFINE_string('dataset', '', 'Which dataset to generate data for')

# Preprocessing config
flags.DEFINE_boolean('output_unigrams', True, 'Whether to output unigrams.')
flags.DEFINE_boolean('output_bigrams', False, 'Whether to output bigrams.')
flags.DEFINE_boolean('output_char', False, 'Whether to output characters.')
flags.DEFINE_boolean('lowercase', True, 'Whether to lowercase document terms.')

# IMDB
flags.DEFINE_string('imdb_input_dir', '', 'The input directory containing the '
                    'IMDB sentiment dataset.')
flags.DEFINE_integer('imdb_validation_pos_start_id', 10621, 'File id of the '
                     'first file in the pos sentiment validation set.')
flags.DEFINE_integer('imdb_validation_neg_start_id', 10625, 'File id of the '
                     'first file in the neg sentiment validation set.')

# DBpedia
flags.DEFINE_string('dbpedia_input_dir', '',
                    'Path to DBpedia directory containing train.csv and '
                    'test.csv.')

# Reuters Corpus (rcv1)
flags.DEFINE_string('rcv1_input_dir', '',
                    'Path to rcv1 directory containing train.csv, unlab.csv, '
                    'and test.csv.')

# Rotten Tomatoes
flags.DEFINE_string('rt_input_dir', '',
                    'The Rotten Tomatoes dataset input directory.')

# The amazon reviews input file to use in either the RT or IMDB datasets.
flags.DEFINE_string('amazon_unlabeled_input_file', '',
                    'The unlabeled Amazon Reviews dataset input file. If set, '
                    'the input file is used to augment RT and IMDB vocab.')

Document = namedtuple('Document',
                      'content is_validation is_test label add_tokens')


def documents(dataset='train',
              include_unlabeled=False,
              include_validation=False):
  """Generates Documents based on FLAGS.dataset.

  Args:
    dataset: str, identifies folder within IMDB data directory, test or train.
    include_unlabeled: bool, whether to include the unsup directory. Only valid
      when dataset=train.
    include_validation: bool, whether to include validation data.

  Yields:
    Document

  Raises:
    ValueError: if include_unlabeled is true but dataset is not 'train'
  """

  if include_unlabeled and dataset != 'train':
    raise ValueError('If include_unlabeled=True, must use train dataset')

  # Set the random seed so that we have the same validation set when running
  # gen_data and gen_vocab.
  random.seed(302)

  ds = FLAGS.dataset
  if ds == 'imdb':
    docs_gen = imdb_documents
  elif ds == 'dbpedia':
    docs_gen = dbpedia_documents
  elif ds == 'rcv1':
    docs_gen = rcv1_documents
  elif ds == 'rt':
    docs_gen = rt_documents
  else:
    raise ValueError('Unrecognized dataset %s' % FLAGS.dataset)

  for doc in docs_gen(dataset, include_unlabeled, include_validation):
    yield doc


def tokens(doc):
  """Given a Document, produces character or word tokens.

  Tokens can be either characters, or word-level tokens (unigrams and/or
  bigrams).

  Args:
    doc: Document to produce tokens from.

  Yields:
    token

  Raises:
    ValueError: if all FLAGS.{output_unigrams, output_bigrams, output_char}
      are False.
  """
  if not (FLAGS.output_unigrams or FLAGS.output_bigrams or FLAGS.output_char):
    raise ValueError(
        'At least one of {FLAGS.output_unigrams, FLAGS.output_bigrams, '
        'FLAGS.output_char} must be true')

  content = doc.content.strip()
  if FLAGS.lowercase:
    content = content.lower()

  if FLAGS.output_char:
    for char in content:
      yield char

  else:
    tokens_ = data_utils.split_by_punct(content)
    for i, token in enumerate(tokens_):
      if FLAGS.output_unigrams:
        yield token

      if FLAGS.output_bigrams:
        previous_token = (tokens_[i - 1] if i > 0 else data_utils.EOS_TOKEN)
        bigram = '_'.join([previous_token, token])
        yield bigram
        if (i + 1) == len(tokens_):
          bigram = '_'.join([token, data_utils.EOS_TOKEN])
          yield bigram


def imdb_documents(dataset='train',
                   include_unlabeled=False,
                   include_validation=False):
  """Generates Documents for IMDB dataset.

  Data from http://ai.stanford.edu/~amaas/data/sentiment/

  Args:
    dataset: str, identifies folder within IMDB data directory, test or train.
    include_unlabeled: bool, whether to include the unsup directory. Only valid
      when dataset=train.
    include_validation: bool, whether to include validation data.

  Yields:
    Document

  Raises:
    ValueError: if FLAGS.imdb_input_dir is empty.
  """
  if not FLAGS.imdb_input_dir:
    raise ValueError('Must provide FLAGS.imdb_input_dir')

  tf.logging.info('Generating IMDB documents...')

  def check_is_validation(filename, class_label):
    if class_label is None:
      return False
    file_idx = int(filename.split('_')[0])
    is_pos_valid = (class_label and
                    file_idx >= FLAGS.imdb_validation_pos_start_id)
    is_neg_valid = (not class_label and
                    file_idx >= FLAGS.imdb_validation_neg_start_id)
    return is_pos_valid or is_neg_valid

  dirs = [(dataset + '/pos', True), (dataset + '/neg', False)]
  if include_unlabeled:
    dirs.append(('train/unsup', None))

  for d, class_label in dirs:
    for filename in os.listdir(os.path.join(FLAGS.imdb_input_dir, d)):
      is_validation = check_is_validation(filename, class_label)
      if is_validation and not include_validation:
        continue

      with open(os.path.join(FLAGS.imdb_input_dir, d, filename), encoding='utf-8') as imdb_f:
        content = imdb_f.read()
      yield Document(
          content=content,
          is_validation=is_validation,
          is_test=False,
          label=class_label,
          add_tokens=True)

  if FLAGS.amazon_unlabeled_input_file and include_unlabeled:
    with open(FLAGS.amazon_unlabeled_input_file, encoding='utf-8') as rt_f:
      for content in rt_f:
        yield Document(
            content=content,
            is_validation=False,
            is_test=False,
            label=None,
            add_tokens=False)


def dbpedia_documents(dataset='train',
                      include_unlabeled=False,
                      include_validation=False):
  """Generates Documents for DBpedia dataset.

  Dataset linked to at https://github.com/zhangxiangxiao/Crepe.

  Args:
    dataset: str, identifies the csv file within the DBpedia data directory,
      test or train.
    include_unlabeled: bool, unused.
    include_validation: bool, whether to include validation data, which is a
      randomly selected 10% of the data.

  Yields:
    Document

  Raises:
    ValueError: if FLAGS.dbpedia_input_dir is empty.
  """
  del include_unlabeled

  if not FLAGS.dbpedia_input_dir:
    raise ValueError('Must provide FLAGS.dbpedia_input_dir')

  tf.logging.info('Generating DBpedia documents...')

  with open(os.path.join(FLAGS.dbpedia_input_dir, dataset + '.csv')) as db_f:
    reader = csv.reader(db_f)
    for row in reader:
      # 10% of the data is randomly held out
      is_validation = random.randint(1, 10) == 1
      if is_validation and not include_validation:
        continue

      content = row[1] + ' ' + row[2]
      yield Document(
          content=content,
          is_validation=is_validation,
          is_test=False,
          label=int(row[0]) - 1,  # Labels should start from 0
          add_tokens=True)


def rcv1_documents(dataset='train',
                   include_unlabeled=True,
                   include_validation=False):
  # pylint:disable=line-too-long
  """Generates Documents for Reuters Corpus (rcv1) dataset.

  Dataset described at
  http://www.ai.mit.edu/projects/jmlr/papers/volume5/lewis04a/lyrl2004_rcv1v2_README.htm

  Args:
    dataset: str, identifies the csv file within the rcv1 data directory.
    include_unlabeled: bool, whether to include the unlab file. Only valid
      when dataset=train.
    include_validation: bool, whether to include validation data, which is a
      randomly selected 10% of the data.

  Yields:
    Document

  Raises:
    ValueError: if FLAGS.rcv1_input_dir is empty.
  """
  # pylint:enable=line-too-long

  if not FLAGS.rcv1_input_dir:
    raise ValueError('Must provide FLAGS.rcv1_input_dir')

  tf.logging.info('Generating rcv1 documents...')

  datasets = [dataset]
  if include_unlabeled:
    if dataset == 'train':
      datasets.append('unlab')
  for dset in datasets:
    with open(os.path.join(FLAGS.rcv1_input_dir, dset + '.csv')) as db_f:
      reader = csv.reader(db_f)
      for row in reader:
        # 10% of the data is randomly held out
        is_validation = random.randint(1, 10) == 1
        if is_validation and not include_validation:
          continue

        content = row[1]
        yield Document(
            content=content,
            is_validation=is_validation,
            is_test=False,
            label=int(row[0]),
            add_tokens=True)


def rt_documents(dataset='train',
                 include_unlabeled=True,
                 include_validation=False):
  # pylint:disable=line-too-long
  """Generates Documents for the Rotten Tomatoes dataset.

  Dataset available at http://www.cs.cornell.edu/people/pabo/movie-review-data/
  In this dataset, amazon reviews are used for the unlabeled data.

  Args:
    dataset: str, identifies the data subdirectory.
    include_unlabeled: bool, whether to include the unlabeled data. Only valid
      when dataset=train.
    include_validation: bool, whether to include validation data, which is a
      randomly selected 10% of the data.

  Yields:
    Document

  Raises:
    ValueError: if FLAGS.rt_input_dir is empty.
  """
  # pylint:enable=line-too-long

  if not FLAGS.rt_input_dir:
    raise ValueError('Must provide FLAGS.rt_input_dir')

  tf.logging.info('Generating rt documents...')

  data_files = []
  input_filenames = os.listdir(FLAGS.rt_input_dir)
  for inp_fname in input_filenames:
    if inp_fname.endswith('.pos'):
      data_files.append((os.path.join(FLAGS.rt_input_dir, inp_fname), True))
    elif inp_fname.endswith('.neg'):
      data_files.append((os.path.join(FLAGS.rt_input_dir, inp_fname), False))
  if include_unlabeled and FLAGS.amazon_unlabeled_input_file:
    data_files.append((FLAGS.amazon_unlabeled_input_file, None))

  for filename, class_label in data_files:
    with open(filename) as rt_f:
      for content in rt_f:
        if class_label is None:
          # Process Amazon Review data for unlabeled dataset
          if content.startswith('review/text'):
            yield Document(
                content=content,
                is_validation=False,
                is_test=False,
                label=None,
                add_tokens=False)
        else:
          # 10% of the data is randomly held out for the validation set and
          # another 10% of it is randomly held out for the test set
          random_int = random.randint(1, 10)
          is_validation = random_int == 1
          is_test = random_int == 2
          if (is_test and dataset != 'test') or (is_validation and
                                                 not include_validation):
            continue

          yield Document(
              content=content,
              is_validation=is_validation,
              is_test=is_test,
              label=class_label,
              add_tokens=True)