from EIVData.csv_dataset import CSVData
from torch.utils.data import random_split

def load_data(seed=0, splitting_part=0.8, normalize=True):
    """
    Loads the yacht hydrodynamics dataset
    :param seed: Seed for splitting and shuffling the data.
    Defaults to 0.
    :param splitting_part: Which fraction of the data to use as training
    data. Defaults to 0.8.
    :normalize: Whether to normalize the data, defaults to True.
    :returns: yacht_trainset, yacht_testset
    """
    yacht_dataset = CSVData('~/SharedData/AI/datasets/yacht_hydrodynamics/yacht_hydrodynamics.data',
            class_name=6,
            shuffle_seed=seed,
            normalize=normalize,
            header=None,
            delimiter=r"\s+")
    dataset_len = len(yacht_dataset)
    train_len = int(dataset_len*splitting_part)
    test_len = dataset_len - train_len
    yacht_trainset, yacht_testset = random_split(yacht_dataset,
            lengths=[train_len, test_len])
    return yacht_trainset, yacht_testset