def import_data(flatten=False, one_khz = False, train_size=270):
train_input , train_target = bci.load(root = './ data_bci', one_khz=one_khz)
print("Original data format: ")
print ( str ( type ( train_input ) ) , train_input.size() )
print ( str ( type ( train_target ) ) , train_target.size() )
test_input , test_target = bci.load(root = './ data_bci', train = False, one_khz=one_khz)
print ( str ( type ( test_input ) ) , test_input.size())
print ( str ( type ( test_target ) ) , test_target.size())
# split traindata into train and validation datasets
val_input, val_target = train_input[train_size:], train_target[train_size:]
train_input, train_target = train_input[:train_size], train_target[:train_size]
# if flatten the "discard" the timeinformation, and just look at every feature at a certain time as a feature
if flatten == True:
train_input = flatten_input_data(train_input)
val_input = flatten_input_data(val_input)
test_input = flatten_input_data(test_input)
train_data = TensorDataset(train_input, train_target)
val_data = TensorDataset(val_input, val_target)
test_data = TensorDataset(test_input, test_target)
# One-hotting targets
labels = train_data.target_tensor
train_data.target_tensor = torch.LongTensor(labels.size(0), 2).zero_().scatter_(1, labels.view(-1, 1), 1)
labels = val_data.target_tensor
val_data.target_tensor = torch.LongTensor(labels.size(0), 2).zero_().scatter_(1, labels.view(-1, 1), 1)
labels = test_data.target_tensor
test_data.target_tensor = torch.LongTensor(labels.size(0), 2).zero_().scatter_(1, labels.view(-1, 1), 1)
print("Modified train_data.data_tensor shape: ", train_data.data_tensor.shape)
print("Modified train_data.target_tensor shape: ", train_data.target_tensor.shape)
print("val_data.data_tensor shape: ", val_data.data_tensor.shape)
return train_data, val_data, test_data
