import json


def json_top_atom_count(json_str):
    """Count the number of atoms in the JSON topology used by wepy HDF5.

    Parameters
    ----------
    json_str : str
        A string of valid JSON in the format of JSON used in WepyHDF5
        and mdtraj HDF5 format.

    Returns
    -------
    n_atoms : int
        Number of atoms in the topology.

    """

    top_d = json.loads(json_str)
    atom_count = 0
    atom_count = 0
    for chain in top_d['chains']:
        for residue in chain['residues']:
            atom_count += len(residue['atoms'])

    return atom_count