from torchvision import transforms
from torch.utils.data import DataLoader




def get_dataset(name, args):

    batch_size = args.batch_size * len(args.gpu.split(','))


    if name == 'brats':
        from dataloaders.BRATS_dataloader_new import Hybrid as BratsDataset
        from dataloaders.BRATS_dataloader_new import RandomPadCrop, ToTensor, RandomFlip
        

        db_train = BratsDataset(split='train', MRIDOWN=args.MRIDOWN, SNR=args.low_field_SNR,
                            transform=transforms.Compose([RandomPadCrop(),  ToTensor()]),  # RandomFlip(),
                            base_dir=args.root_path, input_normalize = args.input_normalize, 
                            use_kspace=args.use_kspace)
        

        db_test  = BratsDataset(split='test', MRIDOWN=args.MRIDOWN, SNR=args.low_field_SNR, 
                            transform=transforms.Compose([ToTensor()]),
                            base_dir=args.root_path, input_normalize = args.input_normalize, 
                            use_kspace=args.use_kspace)



        trainloader    = DataLoader(db_train, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)
        fixtrainloader = DataLoader(db_train, batch_size=1, shuffle=False, num_workers=4, pin_memory=True)
        testloader     = DataLoader(db_test, batch_size=1, shuffle=False, num_workers=4, pin_memory=True)
       
        return trainloader, fixtrainloader, testloader


    else:
        raise NotImplementedError(f'Dataset {name} is not implemented.')