import tensorflow as tf

from resan.utils.nn import bn_dense_layer, dropout, linear
from resan.utils.general import exp_mask_for_high_rank, mask_for_high_rank
from resan.rl_nn import reduce_data_rep_max_len


def reinforced_self_attention(
        rep_tensor, rep_mask, dep_selection, head_selection,
        hn=None, keep_unselected=True,
        scope=None, keep_prob=1., is_train=None, wd=0., activation='elu'
):
    with tf.variable_scope(scope or 'reinforced_self_attention'):
        fw_result = directional_attention_with_selections(
            rep_tensor, rep_mask, dep_selection, head_selection,
            'forward', hn, keep_unselected,
            'forward_resa', keep_prob, is_train, wd, activation
        )
        bw_result = directional_attention_with_selections(
            rep_tensor, rep_mask, dep_selection, head_selection,
            'backward', hn, keep_unselected,
            'backward_resa', keep_prob, is_train, wd, activation
