import os

from tcutility.molecule import load

j = os.path.join

root_dir = j(os.path.split(__file__)[0], "molecules")


def get(name):
    p = j(root_dir, f"{name.removesuffix('.xyz')}.xyz")
    return load(p)


def get_molecules(tags=None):
    for f in os.listdir(root_dir):
        if not f.endswith(".xyz"):
            continue

        mol = get(f)
        if tags is None or any(tag in mol.flags.tags for tag in tags):
            yield mol


if __name__ == "__main__":
    for mol in get_molecules():
        print(mol, mol.flags)
