import torch


def attn_mask_pointing(imgs_len, text_len, seq_types, num_seqs, attn_masking):
    """
    attn_masking options:
    - bottleneck: isolate attention but target images can also look at other images
    - isolate_attn: each image/text can only look at themselves and their corresponding text/images
    - full: full attention across sequences and modalities, no restrictions
    - full_target_query: sequences can attend within themselves (and cross-modality), and target attends everywhere
    - full_target_query_key: same as full_target_query, but now all sequences can also attend to the target
    At the output: The format of the three outputs corresponds to the format of the input of the encoder. It has length
    K. The first token is "[IMG]", that signals the start of the images. Then there are the images for the first
    sequence. Then another token ("[SEP]") that signals the beginning of the images for the second sequence. Then the
    images for the second sequence, and so on. These images do NOT have padding. When the S sequences of images finish,
    there is the token "[TXT]" representing the beginning of the text tokens. Between the last image and the "[TXT]"
    token there is a "[SEP]" token. Then there is the text of the first sequence, then another "[SEP]",
    then the text of the second sequence and so on, until the end. Then the rest until K is padding ("[PAD]" token).
    Between the last text token and the first "[PAD]" token there is a "[TXT]" token too
    - img_locs: len K, True where there are text token embeddings.
    - txt_locs: len K, True where there are image features.
    - attn_mask: tensor(K x K) containing the product of query and key where we limit the attention. Depends on
    `attn_masking` and follows the same structure described above.
    """
    max_num_toks = (imgs_len.sum(dim=1) + text_len.sum(dim=1) + num_seqs * 2).max() + 2
    # all images and text, all separation tokens (2 per seq - one for images, one for text), [IMG] and [TXT]
    # aggregation tokens. This assumes all sequences come with >0 image and text tokens
    attn_mask = torch.zeros((len(imgs_len), max_num_toks, max_num_toks))  # batch size x # queries x # keys
    img_locs = torch.zeros((len(imgs_len), max_num_toks)).bool()
    txt_locs = torch.zeros((len(imgs_len), max_num_toks)).bool()
    attn_mask[:, :1] = 1  # "[IMG]" token can attend everywhere
    if attn_masking == 'full':
        attn_mask[:, :] = 1
    for i, (img_len, txt_len, seq_type, seq_cnt) in enumerate(
            zip(imgs_len, text_len, seq_types, num_seqs)):
        img_start = 1
        txt_start = img_start + sum(img_len) + seq_cnt + 1
        attn_mask[i][txt_start - 1, :] = 1  # [TXT] can attend everywhere
        for j, (ilen, tlen, seq) in enumerate(zip(img_len, txt_len, seq_type)):
            if j == seq_cnt:
                break
            if seq == 0 and not attn_masking == 'isolate_attn':
                if attn_masking == 'bottleneck':
                    # text from target sequence can only attend to its image (and itself)
                    attn_mask[i][txt_start:txt_start + tlen + 1, txt_start:txt_start + tlen + 1] = 1
                    attn_mask[i][txt_start:txt_start + tlen + 1, img_start:img_start + ilen + 1] = 1
                    # img from target sequence can attend to itself, to its text, and all other images
                    attn_mask[i][img_start:img_start + ilen + 1, txt_start:txt_start + tlen + 1] = 1
                    img_start_ = 1
                    for k, (ilen_, _, _) in enumerate(zip(img_len, txt_len, seq_type)):
                        attn_mask[i][img_start:img_start + ilen + 1, img_start_:img_start_ + ilen_ + 1] = 1
                        img_start_ += ilen_ + 1
                # full_target_query and full_target_query_key
                elif attn_masking == 'full_target_query' or attn_masking == 'full_target_query_key':
                    # target sequence can attend to all the other sequences
                    attn_mask[i][img_start:img_start + ilen + 1, :] = 1
                    attn_mask[i][txt_start:txt_start + tlen + 1, :] = 1
                    if attn_masking == 'full_target_query_key':
                        # all other sequences can attend to target sequence
                        attn_mask[i][:, img_start:img_start + ilen + 1] = 1
                        attn_mask[i][:, txt_start:txt_start + tlen + 1] = 1
            else:  # positive and negative sequences can only attend to their inputs (both text and image)
                attn_mask[i][img_start:img_start + ilen + 1, img_start:img_start + ilen + 1] = 1
                attn_mask[i][txt_start:txt_start + tlen + 1, txt_start:txt_start + tlen + 1] = 1
                attn_mask[i][img_start:img_start + ilen + 1, txt_start:txt_start + tlen + 1] = 1
                attn_mask[i][txt_start:txt_start + tlen + 1, img_start:img_start + ilen + 1] = 1
            img_locs[i][img_start:img_start + ilen] = True
            txt_locs[i][txt_start:txt_start + tlen] = True
            img_start += ilen + 1
            txt_start += tlen + 1
        attn_mask[i][:, txt_start:] = 0  # we don't want the original sequence to attend to padding
        attn_mask[i][txt_start:] = 0  # we don't want the original sequence to attend to padding

    return attn_mask, img_locs, txt_locs