import tensorflow as tf


def global_attention(state, hidden_states, encoder, encoder_input_length, scope=None, context=None, **kwargs):
    with tf.variable_scope(scope or 'attention_{}'.format(encoder.name)):
        if context is not None and encoder.use_context:
            state = tf.concat([state, context], axis=1)

        e = compute_energy(hidden_states, state, encoder, input_length=encoder_input_length, **kwargs)
        mask = tf.sequence_mask(encoder_input_length, maxlen=tf.shape(hidden_states)[1], dtype=tf.float32)
        e *= mask

        if encoder.attn_norm_fun == 'none':
            weights = e
        elif encoder.attn_norm_fun == 'sigmoid':
            weights = tf.nn.sigmoid(e)
        elif encoder.attn_norm_fun == 'max':
            weights = tf.one_hot(tf.argmax(e, -1), depth=tf.shape(e)[1])
        else:
            e -= tf.reduce_max(e, axis=1, keep_dims=True)
            T = encoder.attn_temperature or 1.0
            exp = tf.exp(e / T) * mask
            weights = exp / tf.reduce_sum(exp, axis=-1, keep_dims=True)

        weighted_average = tf.reduce_sum(tf.expand_dims(weights, 2) * hidden_states, axis=1)

        return weighted_average, weights


def no_attention(state, hidden_states, *args, **kwargs):
    batch_size = tf.shape(state)[0]
    weighted_average = tf.zeros(shape=tf.stack([batch_size, 0]))
