import torch


def collate_batch(batch: list, use_cuda: bool = False) -> tuple:
    """Collate a batch of samples into tensors for predicting crystal properties.

    Args:
        batch (list): list of tuples for each sample containing
            - atom_fea: Tensor shape [n_i, atom_fea_len]
            - nbr_fea: Tensor shape [n_i, M, nbr_fea_len]
            - neighbor_fea_idx: LongTensor shape [n_i, M]
            - target: Tensor shape [1]
            - material_id: str
        use_cuda (bool, optional): [description]. Defaults to False.

    Returns:
        tuple: contains
            atom_feas: Tensor shape [N, orig_atom_fea_len]
                Atom features from atom type
            nbr_feas: Tensor shape [N, M, nbr_fea_len]
                Bond features of each atom's M neighbors
            neighbor_indices: LongTensor shape [N, M]
                Indices of M neighbors of each atom
            crystal_atom_idx: list of LongTensors of length N0
                Mapping from the crystal idx to atom idx
            target: Tensor shape [N, 1]
                Target value for prediction
            material_ids: list[str]
        where N = sum(n_i); N0 = sum(i)
    """
    batch_atom_indices, batch_neighbor_indices, crystal_atom_idx = [], [], []
    base_idx = 0
    features, targets, *rest = zip(*batch)

    atom_feas, nbr_feas, atom_indices, neighbor_indices = zip(*features)
    for idx, atom_fea in enumerate(atom_feas):

        batch_atom_indices.append(atom_indices[idx] + base_idx)
        batch_neighbor_indices.append(neighbor_indices[idx] + base_idx)

        n_atoms = atom_fea.shape[0]  # number of atoms for this crystal
        crystal_atom_idx.extend([idx] * n_atoms)
        base_idx += n_atoms

    out_features = (
        torch.cat(atom_feas, dim=0),
        torch.cat(nbr_feas, dim=0),
        torch.cat(batch_atom_indices, dim=0),
        torch.cat(batch_neighbor_indices, dim=0),
        torch.LongTensor(crystal_atom_idx),
    )

    targets = torch.cat(targets, dim=0)

    if use_cuda:
        out_features = [tensor.cuda() for tensor in out_features]
        targets = targets.cuda()

    return out_features, targets, *rest