import torch


def batch_training_graph_matrices(matrices,
                                  nc_weight=-1e20, device='cpu'):
    """
    Combine training matrices for a batch of labels.
    Args:
        matrices: list of tuples of matrices for FSTs
        nc_weight - forarded to fst_to_matrices
        device: pytorch device
    Returns:
        The matrices (see fst_to_matrices) for training FSTs (CTC graphs
        that accept only the trainig utterance), concatenated to a large
        padded matrix of size B x N x K, where B is the batch size,
        N maximum number of states, K maximum degree.
    """
    bs = len(matrices)
    max_n = max([m[0].size(0) for m in matrices])
    max_ks = [max([m[i].size(1) for m in matrices])
              for i in range(len(matrices[0]))]
    batched_matrices = []
    for i, m in enumerate(matrices[0]):
        batched_matrices.append(torch.full(
            (bs, max_n, max_ks[i]),
            0 if m.dtype == torch.int64 else nc_weight,
            dtype=m.dtype, device=device
            ))

    for b, ms in enumerate(matrices):
        for i, m in enumerate(ms):
            batched_matrices[i][b, :m.size(0), :m.size(1)] = m
    return batched_matrices