#!/usr/bin/python
# -*- coding:utf-8 -*-
from .mmap_dataset import MMAPDataset
from .dataset_wrapper import MixDatasetWrapper, PretrainDatasetWrapper
from .dataset_prot_sample import ProtSampleDataset


import utils.register as R

def create_dataset(config: dict):
    splits = []
    for split_name in ['train', 'valid', 'test']:
        split_config = config.get(split_name, None)
        if split_config is None:
            splits.append(None)
            continue
        if isinstance(split_config, dict) and 'datasets' in split_config:
            dataset = PretrainDatasetWrapper(split_config['max_n_vertex'],
                *[R.construct(cfg) for cfg in split_config['datasets']]
            )
        else:
            dataset = R.construct(split_config)
        splits.append(dataset)
    return splits  # train/valid/test