#-*-coding:utf-8-*-
import tensorflow as tf
from  tensorflow.contrib.layers import xavier_initializer
def attention(word_embeddings,char_embeddings,dim_word):
    with tf.variable_scope("attention"):
        print("use_attention")
        char_embeddings=tf.layers.dense(char_embeddings,dim_word,activation=tf.tanh,use_bias=False,kernel_initializer=xavier_initializer())
        attention_w = tf.layers.dense(word_embeddings,dim_word, kernel_initializer=xavier_initializer())
        attention_u = tf.layers.dense(char_embeddings, dim_word, kernel_initializer=xavier_initializer())
        attention_output = tf.nn.tanh(attention_w + attention_u)
        attention_output = tf.nn.sigmoid(
            tf.layers.dense(attention_output, dim_word, activation=None, use_bias=False,
                            kernel_initializer=xavier_initializer()))
        word_embeddings = tf.multiply(word_embeddings, attention_output) + tf.multiply((1 - attention_output),
                                                                                       char_embeddings)
        return  word_embeddings

def dot_attention(inputs, memory, hidden,):
    with tf.variable_scope("dot_attention"):
        print("use dot_aatneion")
        inputs=tf.layers.dense(inputs,2*hidden,kernel_initializer=tf.contrib.layers.xavier_initializer())
        outputs=tf.matmul(inputs,tf.transpose(memory,[0,2,1]))
        logits = tf.nn.softmax(outputs)
        outputs = tf.matmul(logits, memory)
        res = tf.concat([inputs, outputs], axis=-1)
        return res

def attentive_attention(inputs,memorys,hidden):
    with tf.variable_scope("attentive_attention"):
        print("use_aatentive_attention")
        attention_W = tf.get_variable(name="attention_W",
                                      shape=[hidden, hidden],
                                      initializer=tf.contrib.layers.xavier_initializer(),
                                      dtype=tf.float32)
        attention_V = tf.get_variable(name="attention_V", shape=[hidden, 1],
                                      initializer=tf.contrib.layers.xavier_initializer(),
                                      dtype=tf.float32)
        attention_U = tf.get_variable(name="attention_U",
                                      shape=[hidden, hidden],
                                      initializer=tf.contrib.layers.xavier_initializer(),
                                      dtype=tf.float32)

        memorys = tf.reshape(memorys, [-1, hidden])

        s = tf.shape(inputs)
        output = tf.reshape(inputs, [-1, hidden])
        atten_hidden = tf.tanh(tf.matmul(memorys, attention_W)+tf.matmul(output, attention_U))
        e_i = tf.matmul(atten_hidden, attention_V)
        e_i = tf.reshape(e_i, [-1, s[1], 1])
        alpha_i = tf.nn.softmax(e_i)
        output = tf.reshape(output, [-1, s[1], hidden])
        c_i = tf.multiply(alpha_i, output)
        new_output = tf.concat([output, c_i], axis=-1)
        return  new_output

def multihead_attention(queries,
                        keys,
                        num_units=None,
                        num_heads=8,
                        dropout=0,
                        causality=False,
                        scope="multihead_attention",
                        reuse=None):
    '''Applies multihead attention.

    Args:
      queries: A 3d tensor with shape of [N, T_q, C_q].
      keys: A 3d tensor with shape of [N, T_k, C_k].
      num_units: A scalar. Attention size.
      dropout_rate: A floating point number.
      is_training: Boolean. Controller of mechanism for dropout.
      causality: Boolean. If true, units that reference the future are masked.
      num_heads: An int. Number of heads.
      scope: Optional scope for `variable_scope`.
      reuse: Boolean, whether to reuse the weights of a previous layer
        by the same name.

    Returns
      A 3d tensor with shape of (N, T_q, C)
    '''
    with tf.variable_scope(scope, reuse=reuse):
        print("use multihead_attention")
        # Set the fall back option for num_units
        if num_units is None:
            num_units = queries.get_shape().as_list()[-1]

        # Linear projections
        Q = tf.layers.dense(queries, num_units, activation=tf.nn.relu)  # (N, T_q, C)
        K = tf.layers.dense(keys, num_units, activation=tf.nn.relu)  # (N, T_k, C)
        V = tf.layers.dense(keys, num_units, activation=tf.nn.relu)  # (N, T_k, C)

        # Split and concat
        Q_ = tf.concat(tf.split(Q, num_heads, axis=2), axis=0)  # (h*N, T_q, C/h)
        K_ = tf.concat(tf.split(K, num_heads, axis=2), axis=0)  # (h*N, T_k, C/h)
        V_ = tf.concat(tf.split(V, num_heads, axis=2), axis=0)  # (h*N, T_k, C/h)

        # Multiplication
        outputs = tf.matmul(Q_, tf.transpose(K_, [0, 2, 1]))  # (h*N, T_q, T_k)

        # Scale
        outputs = outputs / (K_.get_shape().as_list()[-1] ** 0.5)

        # Key Masking
        key_masks = tf.sign(tf.abs(tf.reduce_sum(keys, axis=-1)))  # (N, T_k)
        key_masks = tf.tile(key_masks, [num_heads, 1])  # (h*N, T_k)
        key_masks = tf.tile(tf.expand_dims(key_masks, 1), [1, tf.shape(queries)[1], 1])  # (h*N, T_q, T_k)

        paddings = tf.ones_like(outputs) * (-2 ** 32 + 1)
        outputs = tf.where(tf.equal(key_masks, 0), paddings, outputs)  # (h*N, T_q, T_k)

        # Causality = Future blinding
        if causality:
            diag_vals = tf.ones_like(outputs[0, :, :])  # (T_q, T_k)
            tril = tf.contrib.linalg.LinearOperatorTriL(diag_vals).to_dense()  # (T_q, T_k)
            masks = tf.tile(tf.expand_dims(tril, 0), [tf.shape(outputs)[0], 1, 1])  # (h*N, T_q, T_k)

            paddings = tf.ones_like(masks) * (-2 ** 32 + 1)
            outputs = tf.where(tf.equal(masks, 0), paddings, outputs)  # (h*N, T_q, T_k)

        # Activation
        outputs = tf.nn.softmax(outputs)  # (h*N, T_q, T_k)

        # Query Masking
        query_masks = tf.sign(tf.abs(tf.reduce_sum(queries, axis=-1)))  # (N, T_q)
        query_masks = tf.tile(query_masks, [num_heads, 1])  # (h*N, T_q)
        query_masks = tf.tile(tf.expand_dims(query_masks, -1), [1, 1, tf.shape(keys)[1]])  # (h*N, T_q, T_k)
        outputs *= query_masks  # broadcasting. (N, T_q, C)

        # Dropouts
        outputs = tf.nn.dropout(outputs,dropout)

        # Weighted sum
        outputs = tf.matmul(outputs, V_)  # ( h*N, T_q, C/h)

        # Restore shape
        outputs = tf.concat(tf.split(outputs, num_heads, axis=0), axis=2)  # (N, T_q, C)

        # Residual connection
        outputs += queries

        # Normalize
        outputs = normalize(outputs)  # (N, T_q, C)

    return outputs


def normalize(inputs,
              epsilon=1e-8,
              scope="ln",
              reuse=None):
    '''Applies layer normalization.

    Args:
      inputs: A tensor with 2 or more dimensions, where the first dimension has
        `batch_size`.
      epsilon: A floating number. A very small number for preventing ZeroDivision Error.
      scope: Optional scope for `variable_scope`.
      reuse: Boolean, whether to reuse the weights of a previous layer
        by the same name.

    Returns:
      A tensor with the same shape and data dtype as `inputs`.
    '''
    with tf.variable_scope(scope, reuse=reuse):
        inputs_shape = inputs.get_shape()
        params_shape = inputs_shape[-1:]

        mean, variance = tf.nn.moments(inputs, [-1], keep_dims=True)
        beta = tf.Variable(tf.zeros(params_shape))
        gamma = tf.Variable(tf.ones(params_shape))
        normalized = (inputs - mean) / ((variance + epsilon) ** (.5))
        outputs = gamma * normalized + beta

    return outputs