from rslearn.config import dataset

from olmoearth_run.runner.models.operational_context import OperationalContext
from olmoearth_run.shared.models.config.olmoearth_config import OlmoEarthConfig
from olmoearth_run.runner.tools.config_transpilers.modality_layers import ModalityLayers
from olmoearth_run.runner.tools.config_transpilers.label_output_layers import LabelOutputLayers


def to_dataset_config(olmoearth_config: OlmoEarthConfig, ops_context: OperationalContext) -> dataset.DatasetConfig:
    return dataset.DatasetConfig(
        layers={
            **ModalityLayers.generate_modality_layers(olmoearth_config, ops_context),
            **LabelOutputLayers.generate_label_and_output_layers(olmoearth_config, ops_context)
        },
    )
