import torch
from typing import Tuple
from typing import List


def concat_enc_outs(
    input: torch.LongTensor,
    enc_out: torch.Tensor,
    mask: torch.BoolTensor,
    embedding_size: int,
    padding_idx: int,
) -> Tuple[torch.Tensor, torch.BoolTensor]:
    """
    Concatenate Encoder Outputs.

    Does the whole "FiD" thing; each query/document pair is independently encoded in the
    Encoder, so we need to concatenate all the outputs prior to sending to the decoder.

    :param input:
        [bsz, seqlen] original input to the encoder
    :param enc_out:
        [bsz * n_docs, seqlen] output representations from the encoder
    :param mask:
        encoder mask
    :param embedding_size:
        emb/hidden size of the enc representations
    :param padding_idx:
        pad token index; used for mask purposes.

    :return (new_out, new_mask):
        return the encoder output and encoder mask, appropriately concatenated.
    """
    bsz, n_docs = input.size(0), enc_out.size(0) // input.size(0)
    split_enc_out = enc_out.split([n_docs] * bsz, dim=0)
    split_mask = mask.split([n_docs] * bsz, dim=0)

    concat_outs: List[torch.Tensor] = []
    concat_lengths = []
    for i in range(bsz):
        mask_i = split_mask[i].view(-1)
        out_i = split_enc_out[i].reshape(-1, embedding_size)[mask_i]
        concat_outs.append(out_i)
        concat_lengths.append(out_i.size(0))

    new_out = enc_out.new(bsz, max(concat_lengths), embedding_size)
    new_mask: torch.BoolTensor = mask.new(bsz, max(concat_lengths))  # type: ignore
    new_out.fill_(padding_idx)
    new_mask.fill_(False)

    for i, (out_i, length_i) in enumerate(zip(concat_outs, concat_lengths)):
        new_out[i, :length_i] = out_i
        new_mask[i, :length_i] = True

    return new_out, new_mask