"""Data handling for MoltenFlow."""

from .latents import (
    generate_toy_latents,
    generate_toy_targets,
    generate_toy_conditions,
    get_latents,
    rdkit_fingerprints,
)
from .dataset import (
    load_csv_dataset,
    load_processed_dataset,
    load_property_dataset,
    merge_druglike_datasets,
    MolDataset,
)
from .smiles_dataset import SmilesDataset, batchify, download_dataset
from .transforms import TargetScaler
from .splits import scaffold_split, random_split, load_split_indices, save_split_indices
from .backends import load_data_for_surrogate
from .toy_dataset import (
    pool_latents,
    extract_latents_from_vae,
    generate_toydata_from_latents,
    save_toydata,
    load_toydata,
    create_toydata_from_vae,
)

__all__ = [
    # Latent generation
    "generate_toy_latents",
    "generate_toy_targets",
    "generate_toy_conditions",
    "get_latents",
    "rdkit_fingerprints",
    # Dataset loading
    "load_csv_dataset",
    "load_processed_dataset",
    "load_property_dataset",
    "merge_druglike_datasets",
    "MolDataset",
    # SMILES datasets
    "SmilesDataset",
    "batchify",
    "download_dataset",
    # Transforms
    "TargetScaler",
    # Splits
    "scaffold_split",
    "random_split",
    "load_split_indices",
    "save_split_indices",
    # Backends
    "load_data_for_surrogate",
    # Toy dataset from real latents
    "pool_latents",
    "extract_latents_from_vae",
    "generate_toydata_from_latents",
    "save_toydata",
    "load_toydata",
    "create_toydata_from_vae",
]
