import torch


def get_concept_similarity(phrase_embedding_t, box_class_embedding_t, f_aggregate, f_similarity, f_activation):
    """
    Return the similarity between bounding box's class embedding (i.e., textual representation) and phrase.

    Please note that we mask `phrase_embedding` with `phrase_mask` in order to inhibit contribution from masked words
    computed by `f_aggregate`. However, the tensor should be already masked.

    :param phrase_embedding_t: A ([*, d2, d3, d4], [*, d2, d3, 1]) tuple of tensors
    :param box_class_embedding_t: A ([*, d1, d4], [*, d1, 1]) tuple of tensors
    :param f_aggregate: A function that computes aggregate representation of a phrase
    :param f_similarity: A similarity function
    :param f_activation: An activation function, applied on final similarity score
    :return: A [*, d2, d1] tensor
    """
    box_class_embedding, box_class_embedding_mask = box_class_embedding_t
    phrase_embedding, phrase_embedding_mask = phrase_embedding_t

    n_box = box_class_embedding.size()[-2]
    n_ph = phrase_embedding.size()[-3]

    phrase_embedding = f_aggregate(phrase_embedding_t, box_class_embedding_t, dim=-2, f_similarity=f_similarity)

    box_class_embedding = box_class_embedding.unsqueeze(-3).repeat(1, n_ph, 1, 1)
    phrase_embedding = phrase_embedding.unsqueeze(-2).repeat(1, 1, n_box, 1)

    similarity = f_similarity(box_class_embedding, phrase_embedding, dim=-1)

    phrase_embedding_synthetic_mask = phrase_embedding_mask.squeeze(-1).sum(dim=-1, keepdims=True)
    box_class_embedding_mask = box_class_embedding_mask.squeeze(-1).unsqueeze(-2)

    similarity = torch.masked_fill(similarity, mask=phrase_embedding_synthetic_mask == 0, value=-1)
    similarity = torch.masked_fill(similarity, mask=box_class_embedding_mask == 0, value=-1)

    return f_activation(similarity)