import math


def train_test_split(observations, labels, num_labelled_samples,
                     train_percentage):
  """Splits observations and labels in train and test sets.

  Args:
    observations: Numpy array containing the observations with shape
      (num_labelled_samples, 64, 64, num_channels).
    labels: Numpy array containing the observed factors of variations with shape
      (num_labelled_samples, num_factors).
    num_labelled_samples: How many labelled observations are expected. Used to
      check that the observations have the right shape.
    train_percentage: Float in [0,1] indicating which fraction of the labelled
      data should be used for training.

  Returns:
    observations_train: Numpy array of shape
      (train_percentage * num_labelled_samples, 64, 64, num_channels) containing
      the observations for the training.
    labels_train: Numpy array of shape (train_percentage * num_labelled_samples,
      num_factors) containing the observed factors of variation for the training
      observations.
    observations_test: Numpy array containing the observations for the testing
      with shape ((1-train_percentage) * num_labelled_samples, 64, 64,
        num_channels)
    labels_test: Numpy array containing the observed factors of variation for
      the testing data with shape ((1-train_percentage) * num_labelled_samples,
      num_factors).
  """
  assert observations.shape[0] == num_labelled_samples, \
      "Wrong observations shape."
  num_labelled_samples_train = int(
      math.ceil(num_labelled_samples * train_percentage))
  num_labelled_samples_test = num_labelled_samples - num_labelled_samples_train
  observations_train = observations[:num_labelled_samples_train, :, :, :]
  observations_test = observations[num_labelled_samples_train:, :, :, :]

  labels_train = labels[:num_labelled_samples_train, :]
  labels_test = labels[num_labelled_samples_train:, :]
  assert labels_test.shape[0] == num_labelled_samples_test, "Wrong size test."
  return observations_train, labels_train, observations_test, labels_test