import logging


def check_batch(batch, max_length=None):
    """print data shapes and check max_length"""
    def _print(k, d):
        if k in d:
            print(f'\t{k}: {d[k].shape}')

    logging.info(f'Get protein_name: {batch["name"]}')    
    for k in ['aatype', 'msa_feat', 'extra_msa', 'masked_msa_only']:
        _print(k, batch["feat"])
    for k in ['all_atom_positions']:
        _print(k, batch["label"])

    L = batch["feat"]['aatype'].shape[2]
    if not max_length is None and L > max_length:
        print(f'\tskip {batch["name"]} due to two long length')
        return False
    return True