import mlipx
import zntrack

from models import MODELS

project = zntrack.Project()

slabs = []
{% if slab_config %}
with project.group("initialize"):
    slabs.append(mlipx.BuildASEslab(**{{ slab_config }}).frames)
{% endif %}

adsorbates = []
{% if smiles %}
with project.group("initialize"):
    for smiles in {{ smiles }}:
        adsorbates.append(mlipx.Smiles2Conformers(smiles=smiles, num_confs=1).frames)
{% endif %}

for model_name, model in MODELS.items():
    for idx, slab in enumerate(slabs):
        for jdx, adsorbate in enumerate(adsorbates):
            with project.group(model_name, str(idx)):
                _ = mlipx.RelaxAdsorptionConfigs(
                slabs=slab,
                adsorbates=adsorbate,
                model=model,
            )

project.build()
