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))
{% endif %}{% if material_ids %}
with project.group("initialize"):
    for material_id in {{ material_ids }}:
        frames.append(mlipx.MPRester(search_kwargs={"material_ids": [material_id]}))
{% endif %}{% if smiles %}
with project.group("initialize"):
    for smiles in {{ smiles }}:
        frames.append(mlipx.Smiles2Conformers(smiles=smiles, num_confs=1))
{% endif %}

for model_name, model in MODELS.items():
    for idx, data in enumerate(frames):
        with project.group(model_name, str(idx)):
            pd = mlipx.PourbaixDiagram(data=data.frames, model=model, pH=1.0, V=1.8)

project.build()
