<filename>heartpredictions/LSTM/create_dataloaders.py
import torch

def create_dataloaders(dataset, split_proportions, batch_size, display_informations=False):
    lengths = [round(len(dataset) * split) for split in split_proportions]
    r = 0
    for i in range(len(lengths)):
        r_tmp = lengths[i] % 3
        lengths[i] = lengths[i] - r_tmp
        r += r_tmp
    lengths[2] = lengths[2] + r


    #train_dataset, val_dataset, test_dataset = torch.utils.data.random_split(dataset, lengths=lengths, generator=torch.Generator().manual_seed(42))
    
    train_dataset = torch.utils.data.Subset(dataset, range(0, lengths[0]))
    val_dataset = torch.utils.data.Subset(dataset, range(lengths[0], lengths[0] + lengths[1]))
    test_dataset = torch.utils.data.Subset(dataset, range(lengths[0] + lengths[1], lengths[0] + lengths[1] + lengths[2]))
    print("lengths : ", lengths)
    train_dataloader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=1,
        prefetch_factor=1,
        persistent_workers=False,
        pin_memory=True
    )

    val_dataloader = torch.utils.data.DataLoader(
        val_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=4,
        prefetch_factor=2,
        persistent_workers=False,
        pin_memory=True
    )

    test_dataloader = torch.utils.data.DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=4,
        prefetch_factor=2,
        persistent_workers=False,
        pin_memory=True
    )

    if display_informations:
        print(f'Total dataset: {len(train_dataloader) + len(val_dataloader) + len(test_dataloader)}, '
            f'train dataset: {len(train_dataloader)}, val dataset: {len(val_dataloader)}, test_dataset: {len(test_dataloader)}')
    return train_dataloader, val_dataloader, test_dataloader

