# 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


def TestData(
    lot: str = "wb97x-631gd",
    verbose: bool = True,
    download: bool = True,
    dummy_properties: tp.Optional[tp.Dict[str, tp.Any]] = None,
    skip_check: bool = False,
) -> ANIDataset:
    r"""
    GDB subset, only for debugging and code test purposes
    """
    return _fetch_and_create_builtin_dataset(
        root=datasets_dir(),
        ds_name="TestData",
        lot=lot,
        verbose=verbose,
        download=download,
        dummy_properties=dummy_properties,
        skip_check=skip_check,
        suffix=".h5",
    )


def TestDataIons(
    lot: str = "b973c-def2mtzvp",
    verbose: bool = True,
    download: bool = True,
    dummy_properties: tp.Optional[tp.Dict[str, tp.Any]] = None,
    skip_check: bool = False,
) -> ANIDataset:
    r"""
    Only for debugging and test purposes, includes forces, dipoles and charges
    """
    return _fetch_and_create_builtin_dataset(
        root=datasets_dir(),
        ds_name="TestDataIons",
        lot=lot,
        verbose=verbose,
        download=download,
        dummy_properties=dummy_properties,
        skip_check=skip_check,
        suffix=".h5",
    )


def TestDataForcesDipoles(
    lot: str = "b973c-def2mtzvp",
    verbose: bool = True,
    download: bool = True,
    dummy_properties: tp.Optional[tp.Dict[str, tp.Any]] = None,
    skip_check: bool = False,
) -> ANIDataset:
    r"""
    Only for debugging and code testing purposes, includes forces and dipoles
    """
    return _fetch_and_create_builtin_dataset(
        root=datasets_dir(),
        ds_name="TestDataForcesDipoles",
        lot=lot,
        verbose=verbose,
        download=download,
        dummy_properties=dummy_properties,
        skip_check=skip_check,
        suffix=".h5",
    )


def IonsVeryHeavy(
    lot: str = "b973c-def2mtzvp",
    verbose: bool = True,
    download: bool = True,
    dummy_properties: tp.Optional[tp.Dict[str, tp.Any]] = None,
    skip_check: bool = False,
) -> ANIDataset:
    r"""
    Dataset that includes ions, with H,C,N,O,F,S,Cl elements and at least one of
    Si,As,Br,Se,P,B,I (disjoint from LightIons and IonsHeavy) This dataset is not meant
    to be trained to on its own
    """
    return _fetch_and_create_builtin_dataset(
        root=datasets_dir(),
        ds_name="IonsVeryHeavy",
        lot=lot,
        verbose=verbose,
        download=download,
        dummy_properties=dummy_properties,
        skip_check=skip_check,
        suffix=".h5",
    )


def IonsHeavy(
    lot: str = "b973c-def2mtzvp",
    verbose: bool = True,
    download: bool = True,
    dummy_properties: tp.Optional[tp.Dict[str, tp.Any]] = None,
    skip_check: bool = False,
) -> ANIDataset:
    r"""
    WARNING: This dataset may have incorrect energies and/or forces. Dataset that
    includes ions, with H,C,N,O elements and at least one of F,S,Cl (disjoint from
    IonsLight) This dataset is not meant to be trained to on its own
    """
    return _fetch_and_create_builtin_dataset(
        root=datasets_dir(),
        ds_name="IonsHeavy",
        lot=lot,
        verbose=verbose,
        download=download,
        dummy_properties=dummy_properties,
        skip_check=skip_check,
        suffix=".h5",
    )


def IonsLight(
    lot: str = "b973c-def2mtzvp",
    verbose: bool = True,
    download: bool = True,
    dummy_properties: tp.Optional[tp.Dict[str, tp.Any]] = None,
    skip_check: bool = False,
) -> ANIDataset:
    r"""
    WARNING: This dataset may have incorrect energies and/or forces. Dataset that
    includes ions, with H,C,N,O elements only. Not meant to be trained to on its own
    """
    return _fetch_and_create_builtin_dataset(
        root=datasets_dir(),
        ds_name="IonsLight",
        lot=lot,
        verbose=verbose,
        download=download,
        dummy_properties=dummy_properties,
        skip_check=skip_check,
        suffix=".h5",
    )


def ANI1q(
    lot: str = "wb97x-631gd",
    verbose: bool = True,
    download: bool = True,
    dummy_properties: tp.Optional[tp.Dict[str, tp.Any]] = None,
    skip_check: bool = False,
) -> ANIDataset:
    r"""
    Very limited subset of ANI-1x for which 'atomic CM5 charges' are available. This
    dataset is not meant to be trained to on its own. Originally published in *The
    ANI-1ccx and ANI-1x data sets, coupled-cluster and density functional theory
    properties for molecules*. DOI: '10.1038/s41597-020-0473-z'
    """
    return _fetch_and_create_builtin_dataset(
        root=datasets_dir(),
        ds_name="ANI1q",
        lot=lot,
        verbose=verbose,
        download=download,
        dummy_properties=dummy_properties,
        skip_check=skip_check,
        suffix=".h5",
    )


def ANI2qHeavy(
    lot: str = "wb97x-631gd",
    verbose: bool = True,
    download: bool = True,
    dummy_properties: tp.Optional[tp.Dict[str, tp.Any]] = None,
    skip_check: bool = False,
) -> ANIDataset:
    r"""
    Subset of ANI-2x 'heavy' for which 'atomic CM5 charges' are available. This dataset
    is not meant to be trained to on its own. Originally published in *TODO*. DOI:
    'TODO'
    """
    return _fetch_and_create_builtin_dataset(
        root=datasets_dir(),
        ds_name="ANI2qHeavy",
        lot=lot,
        verbose=verbose,
        download=download,
        dummy_properties=dummy_properties,
        skip_check=skip_check,
        suffix=".h5",
    )


def ANI1ccx(
    lot: str = "ccsd(t)star-cbs",
    verbose: bool = True,
    download: bool = True,
    dummy_properties: tp.Optional[tp.Dict[str, tp.Any]] = None,
    skip_check: bool = False,
) -> ANIDataset:
    r"""
    This dataset also has Hartree Fock (HF) energies, RI-MP2 energies and forces and
    DPLNO-CCSD(T) energies for different basis sets and PNO settings. This dataset was
    originally used for transfer learning, not direct training. Originally published in
    *The ANI-1ccx and ANI-1x data sets, coupled-cluster and density functional theory
    properties for molecules*. DOI: '10.1038/s41597-020-0473-z'
    """
    return _fetch_and_create_builtin_dataset(
        root=datasets_dir(),
        ds_name="ANI1ccx",
        lot=lot,
        verbose=verbose,
        download=download,
        dummy_properties=dummy_properties,
        skip_check=skip_check,
        suffix=".h5",
    )


def ANI1x(
    lot: str = "wb97x-631gd",
    verbose: bool = True,
    download: bool = True,
    dummy_properties: tp.Optional[tp.Dict[str, tp.Any]] = None,
    skip_check: bool = False,
) -> ANIDataset:
    r"""
    Originally published in *The ANI-1ccx and ANI-1x data sets, coupled-cluster and
    density functional theory properties for molecules*. DOI:
    '10.1038/s41597-020-0473-z'
    """
    return _fetch_and_create_builtin_dataset(
        root=datasets_dir(),
        ds_name="ANI1x",
        lot=lot,
        verbose=verbose,
        download=download,
        dummy_properties=dummy_properties,
        skip_check=skip_check,
        suffix=".h5",
    )


def ANI2x(
    lot: str = "wb97x-631gd",
    verbose: bool = True,
    download: bool = True,
    dummy_properties: tp.Optional[tp.Dict[str, tp.Any]] = None,
    skip_check: bool = False,
) -> ANIDataset:
    r"""
    In all cases the 'v2' and '2x' datasets are supersets of the 'v1' and '1x' datasets,
    so everything that is in the v1/1x datasets is also in the v2/2x datasets, which
    contain extra structures, except for some wb97X/def2-TZVPP data points for which
    there are 'v1' values but not 'v2' values. Originally published in *TODO*. DOI:
    'TODO'
    """
    return _fetch_and_create_builtin_dataset(
        root=datasets_dir(),
        ds_name="ANI2x",
        lot=lot,
        verbose=verbose,
        download=download,
        dummy_properties=dummy_properties,
        skip_check=skip_check,
        suffix=".h5",
    )


def COMP6v1(
    lot: str = "wb97x-631gd",
    verbose: bool = True,
    download: bool = True,
    dummy_properties: tp.Optional[tp.Dict[str, tp.Any]] = None,
    skip_check: bool = False,
) -> ANIDataset:
    r"""
    Test set, not meant for direct training. In all cases the 'v2' and '2x' datasets are
    supersets of the 'v1' and '1x' datasets, so everything that is in the v1/1x datasets
    is also in the v2/2x datasets, which contain extra structures, except for some
    wb97X/def2-TZVPP data points for which there are 'v1' values but not 'v2' values.
    Note that the ANI-BenchMD, S66x8 and the ``13`` molecules (with 13 heavy atoms) of
    GDB-10to13 were recalculated using ORCA 5.0 instead of 4.2, with default integration
    grids. The numerical difference is not significant for the purposes of training.
    Originally published in *TODO*. DOI: 'TODO'
    """
    return _fetch_and_create_builtin_dataset(
        root=datasets_dir(),
        ds_name="COMP6v1",
        lot=lot,
        verbose=verbose,
        download=download,
        dummy_properties=dummy_properties,
        skip_check=skip_check,
        suffix=".h5",
    )


def COMP6v2(
    lot: str = "wb97x-631gd",
    verbose: bool = True,
    download: bool = True,
    dummy_properties: tp.Optional[tp.Dict[str, tp.Any]] = None,
    skip_check: bool = False,
) -> ANIDataset:
    r"""
    Test set, not meant for direct training. Note that the ANI-BenchMD, S66x8 and the
    ``13`` molecules (with 13 heavy atoms) of GDB-10to13 were recalculated using ORCA
    5.0 instead of 4.2, with default integration grids. The numerical difference is not
    significant for the purposes of training. Originally published in *TODO*. DOI:
    'TODO'
    """
    return _fetch_and_create_builtin_dataset(
        root=datasets_dir(),
        ds_name="COMP6v2",
        lot=lot,
        verbose=verbose,
        download=download,
        dummy_properties=dummy_properties,
        skip_check=skip_check,
        suffix=".h5",
    )


def ANI1e(
    lot: str = "wb97x-631gd",
    verbose: bool = True,
    download: bool = True,
    dummy_properties: tp.Optional[tp.Dict[str, tp.Any]] = None,
    skip_check: bool = False,
) -> ANIDataset:
    r"""
    Structures corresponding to all smiles extracted from the ANI-1x dataset, embedded
    in 3D space and optimized with PM7. This dataset does not have forces, but it has
    other physical properties: Rotational constants A, B, C (GHz); dipole and
    polarizability magnitudes (Debye and a_0^3 respectively); energy of HOMO and LUMO,
    and HOMO-LUMO gap (Ha); average <r^2> (spatial extent, a_0^2); zero point
    vibrational energies (ZPVE, Ha); zero Kelvin internal energy (Ha); thermal
    quantities, U, H, G, C_v, at 298.15 K (C_v in cal/K/mol, rest in Ha). Originally
    published in *ANI-1E: An equilibrium database from the ANI-1 database*. DOI: 'TODO'
    """
    return _fetch_and_create_builtin_dataset(
        root=datasets_dir(),
        ds_name="ANI1e",
        lot=lot,
        verbose=verbose,
        download=download,
        dummy_properties=dummy_properties,
        skip_check=skip_check,
        suffix=".h5",
    )


class _DatasetId(Enum):
    TESTDATA = "TestData"
    TESTDATAIONS = "TestDataIons"
    TESTDATAFORCESDIPOLES = "TestDataForcesDipoles"
    IONSVERYHEAVY = "IonsVeryHeavy"
    IONSHEAVY = "IonsHeavy"
    IONSLIGHT = "IonsLight"
    ANI1Q = "ANI1q"
    ANI2QHEAVY = "ANI2qHeavy"
    ANI1CCX = "ANI1ccx"
    ANI1X = "ANI1x"
    ANI2X = "ANI2x"
    COMP6V1 = "COMP6v1"
    COMP6V2 = "COMP6v2"
    ANI1E = "ANI1e"


class _LotId(Enum):
    DEFAULT = "default"
    ALL = "all"
    B973C_DEF2MTZVP = "b973c-def2mtzvp"
    CCSD_PTP_STAR_CBS = "ccsd(t)star-cbs"
    WB97MD3BJ_DEF2TZVPP = "wb97md3bj-def2tzvpp"
    WB97MV_DEF2TZVPP = "wb97mv-def2tzvpp"
    WB97X_631GD = "wb97x-631gd"
    WB97X_DEF2TZVPP = "wb97x-def2tzvpp"
