from dataloader.cde_loader.cdeDataloader import get_dataloader as cde_loader
from dataloader.PeMSDdataloader import get_dataloader as normal_loader
from dataloader.DCRNNdataloader import get_dataloader as DCRNN_loader
from dataloader.EXPdataloader import get_dataloader as EXP_loader
from dataloader.cde_loader.cdeDataloader import get_dataloader as nrde_loader


def get_dataloader(config, normalizer, single):
    model_name = config["basic"]["model"]
    match model_name:
        case "STGNCDE":
            return cde_loader(config, normalizer, single)
        case "STGNRDE":
            return nrde_loader(config, normalizer, single)
        case "DCRNN":
            return DCRNN_loader(config, normalizer, single)
        case "EXP":
            return EXP_loader(config, normalizer, single)
        case _:
            return normal_loader(config, normalizer, single)