"""IMDB Dataset module for sentiment analysis."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow as tf

from data.util import OOV_CHAR
from data.util import pad_sentence
from data.util import START_CHAR

NUM_CLASS = 2


def load(vocabulary_size, sentence_length):
  """Returns training and evaluation input for imdb dataset.

  Args:
    vocabulary_size: The number of the most frequent tokens
      to be used from the corpus.
    sentence_length: The number of words in each sentence.
      Longer sentences get cut, shorter ones padded.
  Raises:
    ValueError: if the dataset value is not valid.
  Returns:
    A tuple of length 4, for training and evaluation data,
    each being an numpy array.
  """
  (x_train, y_train), (x_test, y_test) = tf.keras.datasets.imdb.load_data(
      path="imdb.npz",
      num_words=vocabulary_size,
      skip_top=0,
      maxlen=None,
      seed=113,
      start_char=START_CHAR,
      oov_char=OOV_CHAR,
      index_from=OOV_CHAR+1)

  x_train_processed = []
  for sen in x_train:
    sen = pad_sentence(sen, sentence_length)
    x_train_processed.append(np.array(sen))
  x_train_processed = np.array(x_train_processed)

  x_test_processed = []
  for sen in x_test:
    sen = pad_sentence(sen, sentence_length)
    x_test_processed.append(np.array(sen))
  x_test_processed = np.array(x_test_processed)

  return x_train_processed, np.eye(NUM_CLASS)[y_train], \
         x_test_processed, np.eye(NUM_CLASS)[y_test]