# Copyright 2017 The TensorFlow Authors. 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.
# ==============================================================================
"""Penn Treebank RNN model definition compatible with eager execution.

Model similar to
https://github.com/uve/models/tree/master/tutorials/rnn/ptb

Usage: python ./rnn_ptb.py --data-path=<path_to_dataset>

Penn Treebank (PTB) dataset from:
http://www.fit.vutbr.cz/~imikolov/rnnlm/simple-examples.tgz
"""

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

import argparse
import os
import sys
import time

import numpy as np
import tensorflow as tf

from tensorflow.contrib.cudnn_rnn.python.layers import cudnn_rnn
from tensorflow.contrib.eager.python import tfe

layers = tf.keras.layers


class RNN(tf.keras.Model):
  """A static RNN.

  Similar to tf.compat.v1.nn.static_rnn, implemented as a class.
  """

  def __init__(self, hidden_dim, num_layers, keep_ratio):
    super(RNN, self).__init__()
    self.keep_ratio = keep_ratio
    self.cells = tf.contrib.checkpoint.List([
        tf.nn.rnn_cell.BasicLSTMCell(num_units=hidden_dim)
        for _ in range(num_layers)
    ])

  def call(self, input_seq, training):
    batch_size = int(input_seq.shape[1])
    for c in self.cells:
      state = c.zero_state(batch_size, tf.float32)
      outputs = []
      input_seq = tf.unstack(input_seq, num=int(input_seq.shape[0]), axis=0)
      for inp in input_seq:
        output, state = c(inp, state)
        outputs.append(output)

      input_seq = tf.stack(outputs, axis=0)
      if training:
        input_seq = tf.nn.dropout(input_seq, self.keep_ratio)
    # Returning a list instead of a single tensor so that the line:
    # y = self.rnn(y, ...)[0]
    # in PTBModel.call works for both this RNN and CudnnLSTM (which returns a
    # tuple (output, output_states).
    return [input_seq]


class Embedding(layers.Layer):
  """An Embedding layer."""

  def __init__(self, vocab_size, embedding_dim, **kwargs):
    super(Embedding, self).__init__(**kwargs)
    self.vocab_size = vocab_size
    self.embedding_dim = embedding_dim

  def build(self, _):
    self.embedding = self.add_variable(
        "embedding_kernel",
        shape=[self.vocab_size, self.embedding_dim],
        dtype=tf.float32,
        initializer=tf.random_uniform_initializer(-0.1, 0.1),
        trainable=True)

  def call(self, x):
    return tf.nn.embedding_lookup(self.embedding, x)


# pylint: disable=not-callable
class PTBModel(tf.keras.Model):
  """LSTM for word language modeling.

  Model described in:
  (Zaremba, et. al.) Recurrent Neural Network Regularization
  http://arxiv.org/abs/1409.2329

  See also:
  https://github.com/uve/models/tree/master/tutorials/rnn/ptb
  """

  def __init__(self,
               vocab_size,
               embedding_dim,
               hidden_dim,
               num_layers,
               dropout_ratio,
               use_cudnn_rnn=True):
    super(PTBModel, self).__init__()

    self.keep_ratio = 1 - dropout_ratio
    self.use_cudnn_rnn = use_cudnn_rnn
    self.embedding = Embedding(vocab_size, embedding_dim)

    if self.use_cudnn_rnn:
      self.rnn = cudnn_rnn.CudnnLSTM(
          num_layers, hidden_dim, dropout=dropout_ratio)
    else:
      self.rnn = RNN(hidden_dim, num_layers, self.keep_ratio)

    self.linear = layers.Dense(
        vocab_size, kernel_initializer=tf.random_uniform_initializer(-0.1, 0.1))
    self._output_shape = [-1, hidden_dim]

  def call(self, input_seq, training):
    """Run the forward pass of PTBModel.

    Args:
      input_seq: [length, batch] shape int64 tensor.
      training: Is this a training call.
    Returns:
      outputs tensors of inference.
    """
    y = self.embedding(input_seq)
    if training:
      y = tf.nn.dropout(y, self.keep_ratio)
    y = self.rnn(y, training=training)[0]
    return self.linear(tf.reshape(y, self._output_shape))


def clip_gradients(grads_and_vars, clip_ratio):
  gradients, variables = zip(*grads_and_vars)
  clipped, _ = tf.clip_by_global_norm(gradients, clip_ratio)
  return zip(clipped, variables)


def loss_fn(model, inputs, targets, training):
  labels = tf.reshape(targets, [-1])
  outputs = model(inputs, training=training)
  return tf.reduce_mean(
      tf.nn.sparse_softmax_cross_entropy_with_logits(
          labels=labels, logits=outputs))


def _divide_into_batches(data, batch_size):
  """Convert a sequence to a batch of sequences."""
  nbatch = data.shape[0] // batch_size
  data = data[:nbatch * batch_size]
  data = data.reshape(batch_size, -1).transpose()
  return data


def _get_batch(data, i, seq_len):
  slen = min(seq_len, data.shape[0] - 1 - i)
  inputs = data[i:i + slen, :]
  target = data[i + 1:i + 1 + slen, :]
  return tf.constant(inputs), tf.constant(target)


def evaluate(model, data):
  """evaluate an epoch."""
  total_loss = 0.0
  total_batches = 0
  start = time.time()
  for _, i in enumerate(range(0, data.shape[0] - 1, FLAGS.seq_len)):
    inp, target = _get_batch(data, i, FLAGS.seq_len)
    loss = loss_fn(model, inp, target, training=False)
    total_loss += loss.numpy()
    total_batches += 1
  time_in_ms = (time.time() - start) * 1000
  sys.stderr.write("eval loss %.2f (eval took %d ms)\n" %
                   (total_loss / total_batches, time_in_ms))
  return total_loss


def train(model, optimizer, train_data, sequence_length, clip_ratio):
  """training an epoch."""

  def model_loss(inputs, targets):
    return loss_fn(model, inputs, targets, training=True)

  grads = tfe.implicit_gradients(model_loss)

  total_time = 0
  for batch, i in enumerate(range(0, train_data.shape[0] - 1, sequence_length)):
    train_seq, train_target = _get_batch(train_data, i, sequence_length)
    start = time.time()
    optimizer.apply_gradients(
        clip_gradients(grads(train_seq, train_target), clip_ratio))
    total_time += (time.time() - start)
    if batch % 10 == 0:
      time_in_ms = (total_time * 1000) / (batch + 1)
      sys.stderr.write("batch %d: training loss %.2f, avg step time %d ms\n" %
                       (batch, model_loss(train_seq, train_target).numpy(),
                        time_in_ms))


class Datasets(object):
  """Processed form of the Penn Treebank dataset."""

  def __init__(self, path):
    """Load the Penn Treebank dataset.

    Args:
      path: Path to the data/ directory of the dataset from Tomas Mikolov's
        webpage - http://www.fit.vutbr.cz/~imikolov/rnnlm/simple-examples.tgz
    """

    self.word2idx = {}  # string -> integer id
    self.idx2word = []  # integer id -> word string
    # Files represented as a list of integer ids (as opposed to list of string
    # words).
    self.train = self.tokenize(os.path.join(path, "ptb.train.txt"))
    self.valid = self.tokenize(os.path.join(path, "ptb.valid.txt"))

  def vocab_size(self):
    return len(self.idx2word)

  def add(self, word):
    if word not in self.word2idx:
      self.idx2word.append(word)
      self.word2idx[word] = len(self.idx2word) - 1

  def tokenize(self, path):
    """Read text file in path and return a list of integer token ids."""
    tokens = 0
    with tf.gfile.Open(path, "r") as f:
      for line in f:
        words = line.split() + ["<eos>"]
        tokens += len(words)
        for word in words:
          self.add(word)

    # Tokenize file content
    with tf.gfile.Open(path, "r") as f:
      ids = np.zeros(tokens).astype(np.int64)
      token = 0
      for line in f:
        words = line.split() + ["<eos>"]
        for word in words:
          ids[token] = self.word2idx[word]
          token += 1

    return ids


def small_model(use_cudnn_rnn):
  """Returns a PTBModel with a 'small' configuration."""
  return PTBModel(
      vocab_size=10000,
      embedding_dim=200,
      hidden_dim=200,
      num_layers=2,
      dropout_ratio=0.,
      use_cudnn_rnn=use_cudnn_rnn)


def large_model(use_cudnn_rnn):
  """Returns a PTBModel with a 'large' configuration."""
  return PTBModel(
      vocab_size=10000,
      embedding_dim=650,
      hidden_dim=650,
      num_layers=2,
      dropout_ratio=0.5,
      use_cudnn_rnn=use_cudnn_rnn)


def test_model(use_cudnn_rnn):
  """Returns a tiny PTBModel for unit tests."""
  return PTBModel(
      vocab_size=100,
      embedding_dim=20,
      hidden_dim=20,
      num_layers=2,
      dropout_ratio=0.,
      use_cudnn_rnn=use_cudnn_rnn)


def main(_):
  tf.enable_eager_execution()

  if not FLAGS.data_path:
    raise ValueError("Must specify --data-path")
  corpus = Datasets(FLAGS.data_path)
  train_data = _divide_into_batches(corpus.train, FLAGS.batch_size)
  eval_data = _divide_into_batches(corpus.valid, 10)

  have_gpu = tfe.num_gpus() > 0
  use_cudnn_rnn = not FLAGS.no_use_cudnn_rnn and have_gpu

  with tf.device("/device:GPU:0" if have_gpu else None):
    # Make learning_rate a Variable so it can be included in the checkpoint
    # and we can resume training with the last saved learning_rate.
    learning_rate = tf.Variable(20.0, name="learning_rate")
    model = PTBModel(corpus.vocab_size(), FLAGS.embedding_dim,
                     FLAGS.hidden_dim, FLAGS.num_layers, FLAGS.dropout,
                     use_cudnn_rnn)
    optimizer = tf.train.GradientDescentOptimizer(learning_rate)
    checkpoint = tf.train.Checkpoint(
        learning_rate=learning_rate, model=model,
        # GradientDescentOptimizer has no state to checkpoint, but noting it
        # here lets us swap in an optimizer that does.
        optimizer=optimizer)
    # Restore existing variables now (learning_rate), and restore new variables
    # on creation if a checkpoint exists.
    checkpoint.restore(tf.train.latest_checkpoint(FLAGS.logdir))
    sys.stderr.write("learning_rate=%f\n" % learning_rate.numpy())

    best_loss = None
    for _ in range(FLAGS.epoch):
      train(model, optimizer, train_data, FLAGS.seq_len, FLAGS.clip)
      eval_loss = evaluate(model, eval_data)
      if not best_loss or eval_loss < best_loss:
        if FLAGS.logdir:
          checkpoint.save(os.path.join(FLAGS.logdir, "ckpt"))
        best_loss = eval_loss
      else:
        learning_rate.assign(learning_rate / 4.0)
        sys.stderr.write("eval_loss did not reduce in this epoch, "
                         "changing learning rate to %f for the next epoch\n" %
                         learning_rate.numpy())


if __name__ == "__main__":
  parser = argparse.ArgumentParser()
  parser.add_argument(
      "--data-path",
      type=str,
      default="",
      help="Data directory of the Penn Treebank dataset from "
      "http://www.fit.vutbr.cz/~imikolov/rnnlm/simple-examples.tgz")
  parser.add_argument(
      "--logdir", type=str, default="", help="Directory for checkpoint.")
  parser.add_argument("--epoch", type=int, default=20, help="Number of epochs.")
  parser.add_argument("--batch-size", type=int, default=20, help="Batch size.")
  parser.add_argument(
      "--seq-len", type=int, default=35, help="Sequence length.")
  parser.add_argument(
      "--embedding-dim", type=int, default=200, help="Embedding dimension.")
  parser.add_argument(
      "--hidden-dim", type=int, default=200, help="Hidden layer dimension.")
  parser.add_argument(
      "--num-layers", type=int, default=2, help="Number of RNN layers.")
  parser.add_argument(
      "--dropout", type=float, default=0.2, help="Drop out ratio.")
  parser.add_argument(
      "--clip", type=float, default=0.25, help="Gradient clipping ratio.")
  parser.add_argument(
      "--no-use-cudnn-rnn",
      action="store_true",
      default=False,
      help="Disable the fast CuDNN RNN (when no gpu)")

  FLAGS, unparsed = parser.parse_known_args()
  tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)
