from torch.utils.data import ConcatDataset
from musicfm.data_loaders.mertlong import MERTDataset
from musicfm.data_loaders.msd import MSDDataset
from musicfm.data_loaders.fma import FMADataset
from musicfm.data_loaders.genius import GeniusDataset
from musicfm.data_loaders.tency import TencyDataset


def get_dataset(dataset, split, num_samples=-1, input_length_s=29.0):
    if dataset == "mertlong":
        return MERTDataset(
            split=split, num_samples=num_samples, input_length_s=input_length_s
        )

    elif dataset == "msd":
        return MSDDataset(
            split=split, num_samples=num_samples, input_length_s=input_length_s
        )

    elif dataset == "fma":
        return FMADataset(
            split=split, num_samples=num_samples, input_length_s=input_length_s
        )

    elif dataset == "genius":
        return GeniusDataset(
            split=split, num_samples=num_samples, input_length_s=input_length_s
        )

    elif dataset == "concat":
        datasets = [
            MERTDataset(
                split=split, num_samples=num_samples, input_length_s=input_length_s
            ),
            GeniusDataset(
                split=split, num_samples=num_samples, input_length_s=input_length_s
            ),
            MSDDataset(
                split=split, num_samples=num_samples, input_length_s=input_length_s
            ),
            TencyDataset(
                split=split, num_samples=num_samples, input_length_s=input_length_s
            ),
        ]
        return ConcatDataset(datasets)

    else:
        print("%s dataset is not supported yet" % dataset)