# coding=utf-8
# Copyright 2018 The Tensor2Tensor 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.
"""MultiModel."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from six.moves import range  # pylint: disable=redefined-builtin

from tensor2tensor.layers import common_attention
from tensor2tensor.layers import common_hparams
from tensor2tensor.layers import common_layers
from tensor2tensor.layers import modalities
from tensor2tensor.models import slicenet
from tensor2tensor.utils import expert_utils
from tensor2tensor.utils import registry
from tensor2tensor.utils import t2t_model

import tensorflow as tf


def conv_res_step(x, hparams, padding, mask):
  """One step of convolutions and mid-residual."""
  k = (hparams.kernel_height, hparams.kernel_width)
  k2 = (hparams.large_kernel_size, 1)
  dilations_and_kernels1 = [((1, 1), k), ((1, 1), k)]
  dilations_and_kernels2 = [((1, 1), k2), ((4, 4), k2)]
  with tf.variable_scope("conv_res_step"):
    y = common_layers.subseparable_conv_block(
        x,
        hparams.filter_size,
        dilations_and_kernels1,
        padding=padding,
        mask=mask,
        separabilities=0,
        name="residual1")
    y = tf.nn.dropout(y, 1.0 - hparams.dropout)
    return common_layers.subseparable_conv_block(
        y,
        hparams.hidden_size,
        dilations_and_kernels2,
        padding=padding,
        mask=mask,
        separabilities=0,
        name="residual2")


def residual_fn2(x, y, hparams):
  y = tf.nn.dropout(y, 1.0 - hparams.dropout)
  return common_layers.layer_norm(x + y)


def residual_fn3(x, y, z, hparams):
  y = tf.nn.dropout(y, 1.0 - hparams.dropout)
  z = tf.nn.dropout(z, 1.0 - hparams.dropout)
  return common_layers.layer_norm(x + y + z)


def conv_experts(xs, hparams, dp, ps, padding, mask, layer_id):
  """Convolutions + Mixture-of-Experts layer."""
  del layer_id  # Unused.
  train = hparams.mode == tf.estimator.ModeKeys.TRAIN,
  conv_out = dp(conv_res_step, xs, hparams, padding, mask)
  loss = 0.0
  moe_hidden_sizes = [hparams.filter_size]
  expert_fn = expert_utils.ffn_expert_fn(hparams.hidden_size, moe_hidden_sizes,
                                         hparams.hidden_size)
  moe_out, loss = expert_utils.distributed_moe(
      dp,
      ps,
      xs,
      train,
      input_size=hparams.hidden_size,
      expert_fn=expert_fn,
      num_experts=hparams.moe_num_experts,
      k=hparams.moe_k,
      loss_coef=1.0)
  return dp(residual_fn3, xs, moe_out, conv_out, hparams), loss


def prepare_decoder(targets, target_space_emb):
  """Prepare decoder."""
  decoder_self_attention_bias = (
      common_attention.attention_bias_lower_triangle(tf.shape(targets)[1]))
  target_space_emb = tf.reshape(target_space_emb, [1, 1, -1])
  target_space_emb = tf.tile(target_space_emb, [tf.shape(targets)[0], 1, 1])
  decoder_input = common_layers.shift_right_3d(
      targets, pad_value=target_space_emb)
  decoder_input = common_attention.add_timing_signal_1d(decoder_input)
  return (decoder_input, decoder_self_attention_bias)


@registry.register_model
class MultiModel(t2t_model.T2TModel):
  """Model to train on multiple tasks simultaneously."""

  @property
  def use_body_sharded(self):
    return True

  def body_sharded(self, sharded_features):
    train = self._hparams.mode == tf.estimator.ModeKeys.TRAIN
    dp = self._data_parallelism