import tensorflow as tf
from keras.layers import Lambda
from configuration import conf


def mask_layer_by_task(task_input, input_tensor, name=None, return_mask=False):
    mask = tf.expand_dims(task_input, axis=-1)
    mask = tf.tile(mask, multiples=[1, 1, input_tensor.shape[1] // conf.num_tasks])
    mask = tf.keras.layers.Flatten()(mask)
    if name is None:
        out = Lambda(lambda x: x * mask)(input_tensor)
    else:
        out = Lambda(lambda x: x * mask, name=name)(input_tensor)
    if return_mask:
        return out, mask
    else:
        return out
