from typing import Dict
from omegaconf import DictConfig
from datasets.video.base_video import BaseAdvancedVideoDataset
from . import BaseDataModule, ResumableDataModule


def _data_module_cls(_, root_cfg: DictConfig, compatible_datasets: Dict):
    dataset_cls = compatible_datasets[root_cfg.dataset._name]
    data_module_cls = None
    if issubclass(dataset_cls, BaseAdvancedVideoDataset):
        data_module_cls = ResumableDataModule
    else:
        data_module_cls = BaseDataModule
    return data_module_cls(root_cfg, compatible_datasets)