import torch


def merge_list_feat(list_feat, list_batch_cnt):
    """
    Merge a list of feat into a single feat

    Args:
        list_feat (list): each element is (N, C)
        list_batch_cnt (list): each element is (B) - [N0, N1, ...]; sum(Ni) == N

    Returns:
        merge_feat (torch.Tensor): (M, C)
        batch_idx (torch.Tensor): (M)
        merge_batch_cnt (torch.Tensor): (B) - [M0, M1, ...]
    """
    assert isinstance(list_feat, list) and isinstance(list_batch_cnt, list)
    assert len(list_feat) == len(list_batch_cnt)
    batch_size = list_batch_cnt[0].shape[0]
    for bi in range(len(list_batch_cnt)):
        assert list_batch_cnt[bi].shape[0] == batch_size

    merge_feat = []
    merge_batch_idx = []
    for bi in range(batch_size):
        for feat, batch_cnt in zip(list_feat, list_batch_cnt):
            start_idx = torch.sum(batch_cnt[:bi]) if bi > 0 else 0
            end_idx = start_idx + batch_cnt[bi]
            # extract feat in this data sample
            merge_feat.append(feat[start_idx: end_idx])  # (Ni, C)
            merge_batch_idx.append(feat.new_zeros(batch_cnt[bi]).fill_(bi))  # (Ni)

    merge_feat = torch.cat(merge_feat, dim=0)
    merge_batch_idx = torch.cat(merge_batch_idx, dim=0)

    merge_batch_cnt = torch.cat([batch_cnt.unsqueeze(1) for batch_cnt in list_batch_cnt], dim=1).sum(dim=1)
    return merge_feat, merge_batch_idx, merge_batch_cnt