import torch


def get_maximum_similarity_box(phrase_embedding_t, box_class_embedding_t, f_similarity):
    """
    Return the similarity (and the index wrt dimensions `[*, d2, d3]`) of the most similar word wrt bounding box's
    class.

    For each query, we have a [*, n_word, n_box, n_feat] tensor. We first compute similarity on the last dimension,
    which will let us to consider the similarity between a word and the class of a box.

    Then compute the maximum among the last dimension, leading us the most similar box class wrt each word.
    (See sketch below).

    The resulting tensor is returned.

           box1 box2 box3 ...
    word1        x                      word1   x
    word2   y                     ==>   word2   y
    word3             z                 word3   z
    ...

    :param box_class_embedding_t: A ([*, d1, d4], [*, d1, 1]) tuple of tensors
    :param phrase_embedding_t: A ([*, d2, d3, d4], [*, d2, d3, 1]) tuple of tensors
    :param f_similarity: A similarity function
    :return: A (Tensor, LongTensor) tuple with dimensions ([*, d2, d3], [*, d2, d3])
    """
    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]
    n_word = phrase_embedding.size()[-2]

    def expand_phrase_embedding(phrase_embedding):
        phrase_embedding = phrase_embedding.unsqueeze(-2)

        dims = [1] * phrase_embedding.dim()
        dims[-2] = n_box

        return phrase_embedding.repeat(*dims)

    def expand_box_class_embedding(box_class_embedding):
        box_class_embedding = box_class_embedding.unsqueeze(-3).unsqueeze(-3)

        dims = [1] * box_class_embedding.dim()
        dims[-4] = n_ph
        dims[-3] = n_word

        return box_class_embedding.repeat(*dims)

    # the following two tensor will have size [*, d2, d3, d1, d4], which in a real word scenario may translate
    # to [*, n_ph, n_word, n_box, n_feat]
    phrase_embedding = expand_phrase_embedding(phrase_embedding)
    box_class_embedding = expand_box_class_embedding(box_class_embedding)

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

    similarity = f_similarity(phrase_embedding, box_class_embedding, dim=-1)  # [*, d2, d3, d1]

    # we need to apply some masking for padded values such us word, phrase and box
    similarity = torch.masked_fill(similarity, mask=phrase_embedding_mask == 0, value=-1)  # word
    similarity = torch.masked_fill(similarity, mask=phrase_embedding_synthetic_mask == 0, value=-1)  # phrase
    similarity = torch.masked_fill(similarity, mask=box_class_embedding_mask == 0, value=-1)  # box

    # then, for each word, keep the box with maximum similarity
    return torch.max(similarity, dim=-1)