import tensorflow as tf

from src.nn_utils.general import exp_mask_for_high_rank, mask_for_high_rank
from src.nn_utils.integration_func import directional_attention_with_dense
from src.nn_utils.nn import bn_dense_layer, linear


def bi_directional_simple_block_attention(
        rep_tensor, rep_mask, block_len=5, scope=None,
        keep_prob=1., is_train=None, wd=0., activation='elu', hn=None):
    with tf.variable_scope(scope or 'bi_directional_simple_block_attn'):

        fw_attn_res = simple_block_attention(
            rep_tensor, rep_mask, block_len, "forward_attn", "forward",
            keep_prob, is_train, wd, activation, hn)
        bw_attn_res = simple_block_attention(
            rep_tensor, rep_mask, block_len, "backward_attn", "backward",
            keep_prob, is_train, wd, activation, hn)
