# This file is automatically generated by tools/builtin-datasets-codegen.py,
# which uses the jinja template engine, *DO NOT MODIFY THIS FILE* unless you
# know what you are doing.
r"""TorchANI Built-in Datasets

Datasets provided by the TorchANI library, in a format ready to batch and train
on. All datasets importable from this module are also directly importable from
`torchani.datasets`. For all datasets, the corresponding docstrings provide
references and associated DOI when possible that point to the original
publications. Please cite the original publications if you use any of these
datasets.

Note that the conformations present in datasets with different LoT *may be
different*. In many cases the number of conformations is similar, and many of
the conformations match, but due to filters applied to the data, and to
different convergence characteristics of the different LoT used, some
calculations may not properly converge, or may be too costly for some LoT and
thus are not included.
"""
import typing as tp
from enum import Enum

from torchani.paths import datasets_dir
from torchani.datasets._utils import _fetch_and_create_builtin_dataset
from torchani.datasets.anidataset import ANIDataset


{% for ds in datasets %}
def {{ ds.name }}(
    lot: str = "{{ ds.default_lot }}",
    verbose: bool = True,
    download: bool = True,
    dummy_properties: tp.Optional[tp.Dict[str, tp.Any]] = None,
    skip_check: bool = False,
) -> ANIDataset:
    r"""
{{ ds.docstr }}
    """
    return _fetch_and_create_builtin_dataset(
        root=datasets_dir(),
        ds_name="{{ ds.name }}",
        lot=lot,
        verbose=verbose,
        download=download,
        dummy_properties=dummy_properties,
        skip_check=skip_check,
        suffix=".h5",
    )


{% endfor %}
class _DatasetId(Enum):
{% for ds in datasets %}
    {{ ds.name | upper }} = "{{ ds.name }}"
{% endfor %}


class _LotId(Enum):
    DEFAULT = "default"
    ALL = "all"
{% for lot in lots %}
    {{ lot | replace("-", "_") | replace("(", "_p") | replace(")", "p_") | upper }} = "{{ lot }}"
{% endfor %}

