# ******************************************************************************
# Copyright 2017-2018 Intel Corporation
#
# 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.
# ******************************************************************************
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import tensorflow as tf

# pylint: disable=no-name-in-module
from tensorflow.python.framework import tensor_shape
from tensorflow.python.keras.layers import Wrapper
from tensorflow.python.layers.convolutional import Conv1D
from tensorflow.python.ops import variable_scope
from tensorflow.python.keras.engine.base_layer import Layer
from tensorflow.python.eager import context
from tensorflow.python.ops import nn_impl
from tensorflow.python.keras import initializers
from tensorflow.python.keras.engine.base_layer import InputSpec
from tensorflow.python.ops import array_ops
from tensorflow.python.framework import ops


# ***NOTE***: The WeightNorm Class is copied from this PR:
# https://github.com/tensorflow/tensorflow/issues/14070
# Once this becomes part of the official TF release, it will be removed
class WeightNorm(Wrapper):
    """This wrapper reparameterizes a layer by decoupling the weight's
    magnitude and direction. This speeds up convergence by improving the
    conditioning of the optimization problem.

    Weight Normalization: A Simple Reparameterization to Accelerate
    Training of Deep Neural Networks: https://arxiv.org/abs/1602.07868
    Tim Salimans, Diederik P. Kingma (2016)

    WeightNorm wrapper works for keras and tf layers.

    ```python
      net = WeightNorm(tf.keras.layers.Conv2D(2, 2, activation='relu'),
             input_shape=(32, 32, 3), data_init=True)(x)
      net = WeightNorm(tf.keras.layers.Conv2D(16, 5, activation='relu'),
                       data_init=True)
      net = WeightNorm(tf.keras.layers.Dense(120, activation='relu'),
                       data_init=True)(net)
      net = WeightNorm(tf.keras.layers.Dense(n_classes),
                       data_init=True)(net)
    ```

    Arguments:
      layer: a layer instance.
      data_init: If `True` use data dependent variable initialization

    Raises:
      ValueError: If not initialized with a `Layer` instance.
      ValueError: If `Layer` does not contain a `kernel` of weights
      NotImplementedError: If `data_init` is True and running graph execution
    """

    def __init__(self, layer, data_init=False, **kwargs):
        if not isinstance(layer, Layer):
            raise ValueError(
                "Please initialize `WeightNorm` layer with a "
                "`Layer` instance. You passed: {input}".format(input=layer)
            )

        if not context.executing_eagerly() and data_init:
            raise NotImplementedError(
                "Data dependent variable initialization is not available for " "graph execution"
            )

        self.initialized = True
        if data_init:
            self.initialized = False

        self.layer_depth = None
        self.norm_axes = None
        super(WeightNorm, self).__init__(layer, **kwargs)
        self._track_trackable(layer, name="layer")

    def _compute_weights(self):
        """Generate weights by combining the direction of weight vector
        with it's norm"""
        with variable_scope.variable_scope("compute_weights"):
            self.layer.kernel = (
                nn_impl.l2_normalize(self.layer.v, axis=self.norm_axes) * self.layer.g
            )

    def _init_norm(self, weights):
        """Set the norm of the weight vector"""
        from tensorflow.python.ops.linalg_ops import norm

        with variable_scope.variable_scope("init_norm"):
            # pylint: disable=no-member
            flat = array_ops.reshape(weights, [-1, self.layer_depth])
            # pylint: disable=no-member
            return array_ops.reshape(norm(flat, axis=0), (self.layer_depth,))

    def _data_dep_init(self, inputs):
        """Data dependent initialization for eager execution"""
        from tensorflow.python.ops.nn import moments
        from tensorflow.python.ops.math_ops import sqrt

        with variable_scope.variable_scope("data_dep_init"):
            # Generate data dependent init values
            activation = self.layer.activation
            self.layer.activation = None
            x_init = self.layer.call(inputs)
            m_init, v_init = moments(x_init, self.norm_axes)
            scale_init = 1.0 / sqrt(v_init + 1e-10)

        # Assign data dependent init values
        self.layer.g = self.layer.g * scale_init
        self.layer.bias = -1 * m_init * scale_init
        self.layer.activation = activation
        self.initialized = True

    # pylint: disable=signature-differs
    def build(self, input_shape):
        """Build `Layer`"""
        input_shape = tensor_shape.TensorShape(input_shape).as_list()
        self.input_spec = InputSpec(shape=input_shape)

        if not self.layer.built:
            self.layer.build(input_shape)
            self.layer.built = False

            if not hasattr(self.layer, "kernel"):
                raise ValueError(
                    "`WeightNorm` must wrap a layer that" " contains a `kernel` for weights"
                )

            # The kernel's filter or unit dimension is -1
            self.layer_depth = int(self.layer.kernel.shape[-1])
            self.norm_axes = list(range(self.layer.kernel.shape.ndims - 1))

            self.layer.v = self.layer.kernel
            self.layer.g = self.layer.add_variable(
                name="g",
                shape=(self.layer_depth,),
                initializer=initializers.get("ones"),
                dtype=self.layer.kernel.dtype,
                trainable=True,
            )

            with ops.control_dependencies([self.layer.g.assign(self._init_norm(self.layer.v))]):
                self._compute_weights()

            self.layer.built = True

        super(WeightNorm, self).build()
        self.built = True

    # pylint: disable=arguments-differ
    def call(self, inputs):
        """Call `Layer`"""
        if context.executing_eagerly():
            if not self.initialized:
                self._data_dep_init(inputs)
            self._compute_weights()  # Recompute weights for each forward pass

        output = self.layer.call(inputs)
        return output

    def compute_output_shape(self, input_shape):
        return tensor_shape.TensorShape(self.layer.compute_output_shape(input_shape).as_list())


class TCN:
    """
    This class defines core TCN architecture.
    This is only the base class, training strategy is not implemented.
    """

    def __init__(self, max_len, n_features_in, hidden_sizes, kernel_size=7, dropout=0.2):
        """
        To use this class,
            1. Inherit this class
            2. Define the training losses in build_train_graph()
            3. Define the training strategy in run()
            4. After the inherited class object is initialized,
               call build_train_graph followed by run

        Args:
            max_len: Maximum length of sequence
            n_features_in: Number of input features (dimensions)
            hidden_sizes: Number of hidden sizes in each layer of TCN (same for all layers)
            kernel_size: Kernel size of convolution filter (same for all layers)
            dropout: Dropout, fraction of activations to drop
        """
        self.max_len = max_len
        self.n_features_in = n_features_in
        self.hidden_sizes = hidden_sizes
        self.kernel_size = kernel_size
        self.dropout = dropout
        self.n_hidden_layers = len(self.hidden_sizes)
        receptive_field_len = self.calculate_receptive_field()
        if receptive_field_len < self.max_len:
            print(
                "Warning! receptive field of the TCN: "
                "%d is less than the input sequence length: %d."
                % (receptive_field_len, self.max_len)
            )
        else:
            print(
                "Receptive field of the TCN: %d, input sequence length: %d."
                % (receptive_field_len, self.max_len)
            )
        self.layer_activations = []

        # toggle this for train/inference mode
        self.training_mode = tf.placeholder(tf.bool, name="training_mode")

        self.sequence_output = None

    def calculate_receptive_field(self):
        """

        Returns:

        """
        return 1 + 2 * (self.kernel_size - 1) * (2 ** self.n_hidden_layers - 1)

    def build_network_graph(self, x, last_timepoint=False):
        """
        Given the input placeholder x, build the entire TCN graph
        Args:
            x: Input placeholder
            last_timepoint: Whether or not to select only the last timepoint to output

        Returns:
            output of the TCN
        """
        # loop and define multiple residual blocks
        with tf.variable_scope("tcn"):
            for i in range(self.n_hidden_layers):
                dilation_size = 2 ** i
                in_channels = self.n_features_in if i == 0 else self.hidden_sizes[i - 1]
                out_channels = self.hidden_sizes[i]
                with tf.variable_scope("residual_block_" + str(i)):
                    x = self._residual_block(
                        x,
                        in_channels,
                        out_channels,
                        dilation_size,
                        (self.kernel_size - 1) * dilation_size,
                    )
                    x = tf.nn.relu(x)
                self.layer_activations.append(x)
            self.sequence_output = x

            # get outputs
            if not last_timepoint:
                prediction = self.sequence_output
            else:
                # last time point size (batch_size, hidden_sizes_encoder)
                width = self.sequence_output.shape[1].value
                lt = tf.squeeze(
                    tf.slice(self.sequence_output, [0, width - 1, 0], [-1, 1, -1]), axis=1
                )
                prediction = tf.layers.Dense(
                    1,
                    kernel_initializer=tf.initializers.random_normal(0, 0.01),
                    bias_initializer=tf.initializers.random_normal(0, 0.01),
                )(lt)

        return prediction

    def _residual_block(self, x, in_channels, out_channels, dilation, padding):
        """
        Defines the residual block
        Args:
            x: Input tensor to residual block
            in_channels: Number of input features (dimensions)
            out_channels: Number of output features (dimensions)
            dilation: Dilation rate
            padding: Padding value

        Returns:
            Output of residual path
        """
        xin = x
        # define two temporal blocks
        for i in range(2):
            with tf.variable_scope("temporal_block_" + str(i)):
                x = self._temporal_block(x, out_channels, dilation, padding)

        # sidepath
        if in_channels != out_channels:
            x_side = tf.layers.Conv1D(
                filters=out_channels,
                kernel_size=1,
                padding="same",
                strides=1,
                activation=None,
                dilation_rate=1,
                kernel_initializer=tf.initializers.random_normal(0, 0.01),
                bias_initializer=tf.initializers.random_normal(0, 0.01),
            )(xin)
        else:
            x_side = xin

        # combine both
        return tf.add(x, x_side)

    def _temporal_block(self, x, out_channels, dilation, padding):
        """
        Defines the temporal block, which is a dilated causual conv layer,
        followed by relu and dropout
        Args:
            x: Input to temporal block
            out_channels: Number of conv filters
            dilation: dilation rate
            padding: padding value

        Returns:
            Tensor output of temporal block
        """
        # conv layer
        x = self._dilated_causal_conv(x, out_channels, dilation, padding)

        x = tf.nn.relu(x)

        # dropout
        batch_size = tf.shape(x)[0]
        x = tf.layers.dropout(
            x,
            rate=self.dropout,
            noise_shape=[batch_size, 1, out_channels],
            training=self.training_mode,
        )

        return x

    # define model
    def _dilated_causal_conv(self, x, n_filters, dilation, padding):
        """
        Defines dilated causal convolution
        Args:
            x: Input activation
            n_filters: Number of convolution filters
            dilation: Dilation rate
            padding: padding value

        Returns:
            Tensor output of convolution
        """
        input_width = x.shape[1].value
        with tf.variable_scope("dilated_causal_conv"):
            # define dilated convolution layer with left side padding
            x = tf.pad(x, tf.constant([[0, 0], [padding, 0], [0, 0]]), "CONSTANT")
            x = WeightNorm(
                Conv1D(
                    filters=n_filters,
                    kernel_size=self.kernel_size,
                    padding="valid",
                    strides=1,
                    activation=None,
                    dilation_rate=dilation,
                    kernel_initializer=tf.initializers.random_normal(0, 0.01),
                    bias_initializer=tf.initializers.random_normal(0, 0.01),
                )
            )(x)

        assert x.shape[1].value == input_width

        return x

    def build_train_graph(self, *args, **kwargs):
        """
        Placeholder for defining training losses and metrics
        """
        raise NotImplementedError("Error! losses for training must be defined")

    def run(self, *args, **kwargs):
        """
        Placeholder for defining training strategy
        """
        raise NotImplementedError("Error! training routine must be defined")


class CommonLayers:
    """
    Class that contains the common layers for language modeling -
            word embeddings and projection layer
    """

    def __init__(self):
        """
        Initialize class
        """
        self.word_embeddings_tf = None
        self.num_words = None
        self.n_features_in = None

    def define_input_layer(
        self, input_placeholder_tokens, word_embeddings, embeddings_trainable=True
    ):
        """
        Define the input word embedding layer
        Args:
            input_placeholder_tokens: tf.placeholder, input to the model
            word_embeddings: numpy array (optional), to initialize the embeddings with
            embeddings_trainable: boolean, whether or not to train the embedding table

        Returns:
            Embeddings corresponding to the data in input placeholder
        """
        with tf.device("/cpu:0"):
            with tf.variable_scope("embedding_layer", reuse=False):
                if word_embeddings is None:
                    initializer = tf.initializers.random_normal(0, 0.01)
                else:
                    initializer = tf.constant_initializer(word_embeddings)
                self.word_embeddings_tf = tf.get_variable(
                    "embedding_table",
                    shape=[self.num_words, self.n_features_in],
                    initializer=initializer,
                    trainable=embeddings_trainable,
                )

                input_embeddings = tf.nn.embedding_lookup(
                    self.word_embeddings_tf, input_placeholder_tokens
                )
        return input_embeddings

    def define_projection_layer(self, prediction, tied_weights=True):
        """
        Define the output word embedding layer
        Args:
            prediction: tf.tensor, the prediction from the model
            tied_weights: boolean, whether or not to tie weights from the input embedding layer

        Returns:
            Probability distribution over vocabulary
        """
        with tf.device("/cpu:0"):
            if tied_weights:
                # tie projection layer and embedding layer
                with tf.variable_scope("embedding_layer", reuse=tf.AUTO_REUSE):
                    softmax_w = tf.matrix_transpose(self.word_embeddings_tf)
                    softmax_b = tf.get_variable("softmax_b", [self.num_words])
                    _, l, k = prediction.shape.as_list()
                    prediction_reshaped = tf.reshape(prediction, [-1, k])
                    mult_out = tf.nn.bias_add(tf.matmul(prediction_reshaped, softmax_w), softmax_b)
                    projection_out = tf.reshape(mult_out, [-1, l, self.num_words])
            else:
                with tf.variable_scope("projection_layer", reuse=False):
                    projection_out = tf.layers.Dense(self.num_words)(prediction)
        return projection_out
