import mlipx
import zntrack

from models import MODELS

project = zntrack.Project()

frames = []
{% if datapath %}
with project.group("initialize"):
    for path in {{ datapath }}:
        frames.append(mlipx.LoadDataFile(path=path).frames)
{% endif %}{% if material_ids %}
with project.group("initialize"):
    for material_id in {{ material_ids }}:
        frames.append(mlipx.MPRester(search_kwargs={"material_ids": [material_id]}).frames)
{% endif %}{% if smiles %}
with project.group("initialize"):
    for smiles in {{ smiles }}:
        frames.append(mlipx.Smiles2Conformers(smiles=smiles, num_confs=1).frames)
{% endif %}

for model_name, model in MODELS.items():
    with project.group(model_name):
        neb = mlipx.HomonuclearDiatomics(
            elements=[],
            data=sum(frames, []),  # Use all elements from all frames
            model=model,
            n_points=100,
            min_distance=0.5,
            max_distance=2.0,
        )

project.build()
