""" prefix_sums.py
    Prefix sum related dataloaders
    Developed as part of DeepThinking project
    July 2021
"""

import torch
from torch.utils import data
from easy_to_hard_data import PrefixSumDataset

# Ignore statemenst for pylint:
#     Too many branches (R0912), Too many statements (R0915), No member (E1101),
#     Not callable (E1102), Invalid name (C0103), No exception (W0702),
#     Too many local variables (R0914), Missing docstring (C0116, C0115),
#     Unused import (W0611).
# pylint: disable=R0912, R0915, E1101, E1102, C0103, W0702, R0914, C0116, C0115, W0611


def get_dataloaders(train_batch_size, test_batch_size, train_data, test_data, inner_data=None,
                    train_split=0.8, shuffle=True):

    dataset = PrefixSumDataset("./data", num_bits=train_data)
    testset = PrefixSumDataset("./data", num_bits=test_data)

    train_split = int(train_split * len(dataset))

    trainset, valset = torch.utils.data.random_split(dataset,
                                                     [train_split,
                                                      int(len(dataset) - train_split)],
                                                     generator=torch.Generator().manual_seed(42))

    trainloader = data.DataLoader(trainset, num_workers=0, batch_size=train_batch_size,
                                  shuffle=shuffle, drop_last=True)
    testloader = data.DataLoader(testset, num_workers=0, batch_size=test_batch_size,
                                 shuffle=False, drop_last=False)
    valloader = data.DataLoader(valset, num_workers=0, batch_size=test_batch_size,
                                shuffle=False, drop_last=False)
    loaders = {"train": trainloader, "test": testloader, "val": valloader}

    if inner_data:
        for inner_size in inner_data:
            loaders[f"inner_{inner_size}"] = data.DataLoader(PrefixSumDataset("./data",
                                                                              num_bits=inner_size),
                                                             num_workers=0,
                                                             batch_size=train_batch_size,
                                                             shuffle=shuffle,
                                                             drop_last=True)
    return loaders