# 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.
# ==============================================================================
#
# Copyright (c) 2019, NVIDIA CORPORATION. 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.
"""Basic sequence-to-sequence model with dynamic RNN support."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import abc
import collections
import os

import tensorflow as tf
import numpy as np


from tensorflow.python.framework import function
from tensorflow.python.ops import math_ops

import attention_wrapper
import model_helper
import beam_search_decoder
from utils import iterator_utils
from utils import math_utils
from utils import misc_utils as utils
from utils import vocab_utils

utils.check_tensorflow_version()

__all__ = ["BaseModel"]


def create_attention_mechanism(
    num_units, memory, source_sequence_length, dtype=None):
  """Create attention mechanism based on the attention_option."""
  # Mechanism
  attention_mechanism = attention_wrapper.BahdanauAttention(
      num_units,
      memory,
      memory_sequence_length=tf.to_int64(source_sequence_length),
      normalize=True, dtype=dtype)
  return attention_mechanism


class BaseModel(object):
  """Sequence-to-sequence base class.
  """

  def __init__(self, hparams, mode, features, scope=None, extra_args=None):
    """Create the model.

    Args:
      hparams: Hyperparameter configurations.
      mode: TRAIN | EVAL | INFER
      features: a dict of input features.
      scope: scope of the model.
      extra_args: model_helper.ExtraArgs, for passing customizable functions.

    """
    self.hparams = hparams
    # Set params
    self._set_params_initializer(hparams, mode, features, scope, extra_args)

    # Train graph
    res = self.build_graph(hparams, scope=scope)
    self._set_train_or_infer(res, hparams)

  def _set_params_initializer(self,
                              hparams,
                              mode,
                              features,
                              scope,
                              extra_args=None):
    """Set various params for self and initialize."""
    self.mode = mode
    self.src_vocab_size = hparams.src_vocab_size
    self.tgt_vocab_size = hparams.tgt_vocab_size
    self.features = features
    self.time_major = hparams.time_major

    if hparams.use_char_encode:
      assert (not self.time_major), ("Can't use time major for"
                                     " char-level inputs.")

    self.dtype = tf.float16 if hparams.use_fp16 else tf.float32

    # extra_args: to make it flexible for adding external customizable code
    self.single_cell_fn = None
    if extra_args:
      self.single_cell_fn = extra_args.single_cell_fn

    # Set num units
    self.num_units = hparams.num_units
    # Set num layers
    self.num_encoder_layers = hparams.num_encoder_layers
    self.num_decoder_layers = hparams.num_decoder_layers
    assert self.num_encoder_layers
    assert self.num_decoder_layers

    # Set num residual layers
    if hasattr(hparams, "num_residual_layers"):  # compatible common_test_utils
      self.num_encoder_residual_layers = hparams.num_residual_layers
      self.num_decoder_residual_layers = hparams.num_residual_layers
    else:
      self.num_encoder_residual_layers = hparams.num_encoder_residual_layers
      self.num_decoder_residual_layers = hparams.num_decoder_residual_layers

    # Batch size
    self.batch_size = tf.size(self.features["source_sequence_length"])

    # Global step
    global_step = tf.train.get_global_step()
    if global_step is not None:
      utils.print_out("global_step already created!")

    self.global_step = tf.train.get_or_create_global_step()
    utils.print_out("model.global_step.name: %s" % self.global_step.name)

    # Initializer
    self.random_seed = hparams.random_seed
    initializer = model_helper.get_initializer(
        hparams.init_op, self.random_seed, hparams.init_weight)
    tf.get_variable_scope().set_initializer(initializer)

    # Embeddings
    self.encoder_emb_lookup_fn = tf.nn.embedding_lookup
    self.init_embeddings(hparams, scope)

  def _set_train_or_infer(self, res, hparams):
    """Set up training."""
    loss = res[1]
    if self.mode == tf.contrib.learn.ModeKeys.TRAIN:
      self.train_loss = loss
      self.word_count = tf.reduce_sum(
          self.features["source_sequence_length"]) + tf.reduce_sum(
              self.features["target_sequence_length"])
    elif self.mode == tf.contrib.learn.ModeKeys.EVAL:
      self.eval_loss = loss
    elif self.mode == tf.contrib.learn.ModeKeys.INFER:
      self.infer_logits = res[0]
      self.infer_loss = loss
      self.sample_id = res[2]

    if self.mode != tf.contrib.learn.ModeKeys.INFER:
      ## Count the number of predicted words for compute ppl.
      self.predict_count = tf.reduce_sum(
          self.features["target_sequence_length"])

    # Gradients and SGD update operation for training the model.
    # Arrange for the embedding vars to appear at the beginning.
    # Only build bprop if running on GPU and using dist_strategy, in which
    # case learning rate, grads and train_op are created in estimator model
    # function.
    with tf.name_scope("learning_rate"):
      self.learning_rate = tf.constant(hparams.learning_rate)
      # warm-up
      self.learning_rate = self._get_learning_rate_warmup(hparams)
      # decay
      self.learning_rate = self._get_learning_rate_decay(hparams)

    if (hparams.use_dist_strategy and
        self.mode == tf.contrib.learn.ModeKeys.TRAIN):
      # Gradients
      params = tf.trainable_variables()
      # Print trainable variables
      utils.print_out("# Trainable variables")
      utils.print_out(
          "Format: <name>, <shape>, <dtype>, <(soft) device placement>")
      for param in params:
        utils.print_out(
            "  %s, %s, %s, %s" % (param.name, str(param.get_shape()),
                                  param.dtype.name, param.op.device))
      utils.print_out("Total params size: %.2f GB" % (4. * np.sum([
          p.get_shape().num_elements()
          for p in params
          if p.shape.is_fully_defined()
      ]) / 2**30))

      # Optimizer
      if hparams.optimizer == "sgd":
        opt = tf.train.GradientDescentOptimizer(self.learning_rate)
      elif hparams.optimizer == "adam":
        opt = tf.train.AdamOptimizer(self.learning_rate)
      else:
        raise ValueError("Unknown optimizer type %s" % hparams.optimizer)
      assert opt is not None

      grads_and_vars = opt.compute_gradients(
          self.train_loss,
          params,
          colocate_gradients_with_ops=hparams.colocate_gradients_with_ops)
      gradients = [x for (x, _) in grads_and_vars]

      clipped_grads, grad_norm = model_helper.gradient_clip(
          gradients, max_gradient_norm=hparams.max_gradient_norm)
      self.grad_norm = grad_norm
      self.params = params
      self.grads = clipped_grads

      self.update = opt.apply_gradients(
          list(zip(clipped_grads, params)), global_step=self.global_step)
    else:
      self.grad_norm = None
      self.update = None
      self.params = None
      self.grads = None

  def _get_learning_rate_warmup(self, hparams):
    """Get learning rate warmup."""
    warmup_steps = hparams.warmup_steps
    warmup_scheme = hparams.warmup_scheme
    utils.print_out("  learning_rate=%g, warmup_steps=%d, warmup_scheme=%s" %
                    (hparams.learning_rate, warmup_steps, warmup_scheme))
    if not warmup_scheme:
      return self.learning_rate

    # Apply inverse decay if global steps less than warmup steps.
    # Inspired by https://arxiv.org/pdf/1706.03762.pdf (Section 5.3)
    # When step < warmup_steps,
    #   learing_rate *= warmup_factor ** (warmup_steps - step)
    if warmup_scheme == "t2t":
      # 0.01^(1/warmup_steps): we start with a lr, 100 times smaller
      warmup_factor = tf.exp(tf.log(0.01) / warmup_steps)
      inv_decay = warmup_factor**(tf.to_float(warmup_steps - self.global_step))
    else:
      raise ValueError("Unknown warmup scheme %s" % warmup_scheme)

    return tf.cond(
        self.global_step < hparams.warmup_steps,
        lambda: inv_decay * self.learning_rate,
        lambda: self.learning_rate,
        name="learning_rate_warump_cond")

  def _get_decay_info(self, hparams):
    """Return decay info based on decay_scheme."""
    if hparams.decay_scheme in [
        "luong5", "luong10", "luong234", "jamesqin1616"
    ]:
      epoch_size, _, _ = iterator_utils.get_effective_epoch_size(hparams)
      num_train_steps = int(hparams.max_train_epochs * epoch_size / hparams.batch_size)
      decay_factor = 0.5
      if hparams.decay_scheme == "luong5":
        start_decay_step = int(num_train_steps / 2)
        decay_times = 5
        remain_steps = num_train_steps - start_decay_step
      elif hparams.decay_scheme == "luong10":
        start_decay_step = int(num_train_steps / 2)
        decay_times = 10
        remain_steps = num_train_steps - start_decay_step
      elif hparams.decay_scheme == "luong234":
        start_decay_step = int(num_train_steps * 2 / 3)
        decay_times = 4
        remain_steps = num_train_steps - start_decay_step
      elif hparams.decay_scheme == "jamesqin1616":
        # dehao@ reported TPU setting max_epoch = 2 and use luong234.
        # They start decay after 2 * 2/3 epochs for 4 times.
        # If keep max_epochs = 8 then decay should start at 8 * 2/(3 * 4) epochs
        # and for (4 *4 = 16) times.
        decay_times = 16
        start_decay_step = int(num_train_steps / 16.)
        remain_steps = num_train_steps - start_decay_step
      decay_steps = int(remain_steps / decay_times)
    elif not hparams.decay_scheme:  # no decay
      start_decay_step = num_train_steps
      decay_steps = 0
      decay_factor = 1.0
    elif hparams.decay_scheme:
      raise ValueError("Unknown decay scheme %s" % hparams.decay_scheme)
    return start_decay_step, decay_steps, decay_factor

  def _get_learning_rate_decay(self, hparams):
    """Get learning rate decay."""
    start_decay_step, decay_steps, decay_factor = self._get_decay_info(hparams)
    utils.print_out("  decay_scheme=%s, start_decay_step=%d, decay_steps %d, "
                    "decay_factor %g" % (hparams.decay_scheme, start_decay_step,
                                         decay_steps, decay_factor))

    return tf.cond(
        self.global_step < start_decay_step,
        lambda: self.learning_rate,
        lambda: tf.train.exponential_decay(  # pylint: disable=g-long-lambda
            self.learning_rate,
            (self.global_step - start_decay_step),
            decay_steps, decay_factor, staircase=True),
        name="learning_rate_decay_cond")

  def init_embeddings(self, hparams, scope):
    """Init embeddings."""
    self.embedding_encoder, self.embedding_decoder = (
        model_helper.create_emb_for_encoder_and_decoder(
            share_vocab=hparams.share_vocab,
            src_vocab_size=self.src_vocab_size,
            tgt_vocab_size=self.tgt_vocab_size,
            src_embed_size=self.num_units,
            tgt_embed_size=self.num_units,
            dtype=self.dtype,
            num_enc_partitions=hparams.num_enc_emb_partitions,
            num_dec_partitions=hparams.num_dec_emb_partitions,
            src_vocab_file=hparams.src_vocab_file,
            tgt_vocab_file=hparams.tgt_vocab_file,
            src_embed_file=hparams.src_embed_file,
            tgt_embed_file=hparams.tgt_embed_file,
            use_char_encode=hparams.use_char_encode,
            scope=scope,
        ))

  def build_graph(self, hparams, scope=None):
    """Subclass must implement this method.

    Creates a sequence-to-sequence model with dynamic RNN decoder API.
    Args:
      hparams: Hyperparameter configurations.
      scope: VariableScope for the created subgraph; default "dynamic_seq2seq".

    Returns:
      A tuple of the form (logits, loss_tuple, final_context_state, sample_id),
      where:
        logits: float32 Tensor [batch_size x num_decoder_symbols].
        loss: loss = the total loss / batch_size.
        final_context_state: the final state of decoder RNN.
        sample_id: sampling indices.

    Raises:
      ValueError: if encoder_type differs from mono and bi, or
        attention_option is not (luong | scaled_luong |
        bahdanau | normed_bahdanau).
    """
    utils.print_out("# Creating %s graph ..." % self.mode)

    # Projection
    with tf.variable_scope(scope or "build_network"):
      with tf.variable_scope("decoder/output_projection"):
        self.output_layer = tf.layers.Dense(
            self.tgt_vocab_size, use_bias=False, name="output_projection",
            dtype=self.dtype)

    with tf.variable_scope(scope or "dynamic_seq2seq", dtype=self.dtype):
      # Encoder
      if hparams.language_model:  # no encoder for language modeling
        utils.print_out("  language modeling: no encoder")
        self.encoder_outputs = None
        encoder_state = None
      else:
        self.encoder_outputs, encoder_state = self._build_encoder(hparams)

      ## Decoder
      logits, sample_id = (
          self._build_decoder(self.encoder_outputs, encoder_state, hparams))

      ## Loss
      if self.mode != tf.contrib.learn.ModeKeys.INFER:
        loss = self._compute_loss(logits, hparams.label_smoothing)
      else:
        loss = tf.constant(0.0)

    return logits, loss, sample_id

  @abc.abstractmethod
  def _build_encoder(self, hparams):
    """Subclass must implement this.

    Build and run an RNN encoder.

    Args:
      hparams: Hyperparameters configurations.

    Returns:
      A tuple of encoder_outputs and encoder_state.
    """
    pass

  def _get_infer_maximum_iterations(self, hparams, source_sequence_length):
    """Maximum decoding steps at inference time."""
    if hparams.tgt_max_len_infer:
      maximum_iterations = hparams.tgt_max_len_infer
      utils.print_out("  decoding maximum_iterations %d" % maximum_iterations)
    else:
      # TODO(thangluong): add decoding_length_factor flag
      decoding_length_factor = 2.0
      max_encoder_length = tf.reduce_max(source_sequence_length)
      maximum_iterations = tf.to_int32(
          tf.round(tf.to_float(max_encoder_length) * decoding_length_factor))
    return maximum_iterations

  def _build_decoder(self, encoder_outputs, encoder_state, hparams):
    """Build and run a RNN decoder with a final projection layer.

    Args:
      encoder_outputs: The outputs of encoder for every time step.
      encoder_state: The final state of the encoder.
      hparams: The Hyperparameters configurations.

    Returns:
      A tuple of final logits and final decoder state:
        logits: size [time, batch_size, vocab_size] when time_major=True.
    """

    ## Decoder.
    with tf.variable_scope("decoder") as decoder_scope:

      ## Train or eval
      if self.mode != tf.contrib.learn.ModeKeys.INFER:
        # [batch, time]
        target_input = self.features["target_input"]
        if self.time_major:
          # If using time_major mode, then target_input should be [time, batch]
          # then the decoder_emb_inp would be [time, batch, dim]
          target_input = tf.transpose(target_input)
        decoder_emb_inp = tf.cast(
            tf.nn.embedding_lookup(self.embedding_decoder, target_input),
            self.dtype)

        if not hparams.use_fused_lstm_dec:
          cell, decoder_initial_state = self._build_decoder_cell(
              hparams, encoder_outputs, encoder_state,
              self.features["source_sequence_length"])

          if hparams.use_dynamic_rnn:
            final_rnn_outputs, _ = tf.nn.dynamic_rnn(
                cell,
                decoder_emb_inp,
                sequence_length=self.features["target_sequence_length"],
                initial_state=decoder_initial_state,
                dtype=self.dtype,
                scope=decoder_scope,
                parallel_iterations=hparams.parallel_iterations,
                time_major=self.time_major)
          else:
            final_rnn_outputs, _ = tf.contrib.recurrent.functional_rnn(
                cell,
                decoder_emb_inp,
                sequence_length=tf.to_int32(
                    self.features["target_sequence_length"]),
                initial_state=decoder_initial_state,
                dtype=self.dtype,
                scope=decoder_scope,
                time_major=self.time_major,
                use_tpu=False)
        else:
          if hparams.pass_hidden_state:
            decoder_initial_state = encoder_state
          else:
            decoder_initial_state = tuple((tf.nn.rnn_cell.LSTMStateTuple(
              tf.zeros_like(s[0]), tf.zeros_like(s[1])) for s in encoder_state))
          final_rnn_outputs = self._build_decoder_fused_for_training(
              encoder_outputs, decoder_initial_state, decoder_emb_inp, self.hparams)

        # We chose to apply the output_layer to all timesteps for speed:
        #   10% improvements for small models & 20% for larger ones.
        # If memory is a concern, we should apply output_layer per timestep.
        logits = self.output_layer(final_rnn_outputs)
        sample_id = None
      ## Inference
      else:
        cell, decoder_initial_state = self._build_decoder_cell(
            hparams, encoder_outputs, encoder_state,
            self.features["source_sequence_length"])

        assert hparams.infer_mode == "beam_search"
        _, tgt_vocab_table = vocab_utils.create_vocab_tables(
            hparams.src_vocab_file, hparams.tgt_vocab_file, hparams.share_vocab)
        tgt_sos_id = tf.cast(
            tgt_vocab_table.lookup(tf.constant(hparams.sos)), tf.int32)
        tgt_eos_id = tf.cast(
            tgt_vocab_table.lookup(tf.constant(hparams.eos)), tf.int32)
        start_tokens = tf.fill([self.batch_size], tgt_sos_id)
        end_token = tgt_eos_id
        beam_width = hparams.beam_width
        length_penalty_weight = hparams.length_penalty_weight
        coverage_penalty_weight = hparams.coverage_penalty_weight

        my_decoder = beam_search_decoder.BeamSearchDecoder(
            cell=cell,
            embedding=self.embedding_decoder,
            start_tokens=start_tokens,
            end_token=end_token,
            initial_state=decoder_initial_state,
            beam_width=beam_width,
            output_layer=self.output_layer,
            length_penalty_weight=length_penalty_weight,
            coverage_penalty_weight=coverage_penalty_weight)

        # maximum_iteration: The maximum decoding steps.
        maximum_iterations = self._get_infer_maximum_iterations(
            hparams, self.features["source_sequence_length"])

        # Dynamic decoding
        outputs, _, _ = tf.contrib.seq2seq.dynamic_decode(
            my_decoder,
            maximum_iterations=maximum_iterations,
            output_time_major=self.time_major,
            swap_memory=True,
            scope=decoder_scope)

        logits = tf.no_op()
        sample_id = outputs.predicted_ids

    return logits, sample_id

  def get_max_time(self, tensor):
    time_axis = 0 if self.time_major else 1
    return tensor.shape[time_axis].value or tf.shape(tensor)[time_axis]

  @abc.abstractmethod
  def _build_decoder_cell(self, hparams, encoder_outputs, encoder_state,
                          source_sequence_length):
    """Subclass must implement this.

    Args:
      hparams: Hyperparameters configurations.
      encoder_outputs: The outputs of encoder for every time step.
      encoder_state: The final state of the encoder.
      source_sequence_length: sequence length of encoder_outputs.

    Returns:
      A tuple of a multi-layer RNN cell used by decoder and the initial state of
      the decoder RNN.
    """
    pass

  def _softmax_cross_entropy_loss(self, logits, labels, label_smoothing):
    """Compute softmax loss or sampled softmax loss."""
    use_defun = os.environ["use_defun"] == "true"
    use_xla = os.environ["use_xla"] == "true"

    # @function.Defun(noinline=True, compiled=use_xla)
    def ComputePositiveCrossent(labels, logits):
      crossent = math_utils.sparse_softmax_crossent_with_logits(
          labels=labels, logits=logits)
      return crossent
    crossent = ComputePositiveCrossent(labels, logits)
    assert crossent.dtype == tf.float32

    def _safe_shape_div(x, y):
      """Divides `x / y` assuming `x, y >= 0`, treating `0 / 0 = 0`."""
      return x // tf.maximum(y, 1)

    @function.Defun(tf.float32, tf.float32, compiled=use_xla)
    def ReduceSumGrad(x, grad):
      """docstring."""
      input_shape = tf.shape(x)
      # TODO(apassos) remove this once device placement for eager ops makes more
      # sense.
      with tf.colocate_with(input_shape):
        output_shape_kept_dims = math_ops.reduced_shape(input_shape, -1)
        tile_scaling = _safe_shape_div(input_shape, output_shape_kept_dims)
      grad = tf.reshape(grad, output_shape_kept_dims)
      return tf.tile(grad, tile_scaling)

    def ReduceSum(x):
      """docstring."""
      return tf.reduce_sum(x, axis=-1)
    if use_defun:
      ReduceSum = function.Defun(
          tf.float32,
          compiled=use_xla,
          noinline=True,
          grad_func=ReduceSumGrad)(ReduceSum)

    if abs(label_smoothing) > 1e-3:
      # pylint:disable=invalid-name
      def ComputeNegativeCrossentFwd(logits):
        """docstring."""
        # [time, batch, dim]
        # [time, batch]
        max_logits = tf.reduce_max(logits, axis=-1)
        # [time, batch, dim]
        shifted_logits = logits - tf.expand_dims(max_logits, axis=-1)
        # Always compute loss in fp32
        shifted_logits = tf.to_float(shifted_logits)
        # [time, batch]
        log_sum_exp = tf.log(ReduceSum(tf.exp(shifted_logits)))
        # [time, batch, dim] - [time, batch, 1] --> reduce_sum(-1) -->
        # [time, batch]
        neg_crossent = ReduceSum(
            shifted_logits - tf.expand_dims(log_sum_exp, axis=-1))
        return neg_crossent

      def ComputeNegativeCrossent(logits):
        return ComputeNegativeCrossentFwd(logits)

      if use_defun:
        ComputeNegativeCrossent = function.Defun(
            compiled=use_xla)(ComputeNegativeCrossent)

      neg_crossent = ComputeNegativeCrossent(logits)
      neg_crossent = tf.to_float(neg_crossent)
      num_labels = logits.shape[-1].value
      crossent = (1.0 - label_smoothing) * crossent - (
          label_smoothing / tf.to_float(num_labels) * neg_crossent)
      # pylint:enable=invalid-name

    return crossent

  def _compute_loss(self, logits, label_smoothing):
    """Compute optimization loss."""
    target_output = self.features["target_output"]
    if self.time_major:
      target_output = tf.transpose(target_output)
    max_time = self.get_max_time(target_output)
    self.batch_seq_len = max_time

    crossent = self._softmax_cross_entropy_loss(
        logits, target_output, label_smoothing)
    assert crossent.dtype == tf.float32

    target_weights = tf.sequence_mask(
        self.features["target_sequence_length"], max_time, dtype=crossent.dtype)
    if self.time_major:
      # [time, batch] if time_major, since the crossent is [time, batch] in this
      # case.
      target_weights = tf.transpose(target_weights)

    loss = tf.reduce_sum(crossent * target_weights) / tf.to_float(
        self.batch_size)

    return loss

  def build_encoder_states(self, include_embeddings=False):
    """Stack encoder states and return tensor [batch, length, layer, size]."""
    assert self.mode == tf.contrib.learn.ModeKeys.INFER
    if include_embeddings:
      stack_state_list = tf.stack(
          [self.encoder_emb_inp] + self.encoder_state_list, 2)
    else:
      stack_state_list = tf.stack(self.encoder_state_list, 2)

    # transform from [length, batch, ...] -> [batch, length, ...]
    if self.time_major:
      stack_state_list = tf.transpose(stack_state_list, [1, 0, 2, 3])

    return stack_state_list