class HyperPlaneDataset(torch.utils.data.Dataset):

  # Warning: read using mutable obects for default input arguments in python.
  def __init__(
      self,
      space_range=[
          -0.75,  # Start
          0.75,  # end
      ],  # ND hypercube
      intercept_range=[-0.5, 0.5],
      normal_offset=[1, 0],
      space_offset=[-2, 2],
      max_normal_angle_degree=180,
      dataset_size=1000,
      noise_density=100,
      pattern_density=100,
      quantization_size=0.01,
      random_seed=0,
      dimension=2):

    assert len(normal_offset) == dimension

    self.max_normal_angle_degree = max_normal_angle_degree
    self.intercept_slope = intercept_range[1] - intercept_range[0]
    self.intercept_bias = intercept_range[0]
    self.normal_offset = normal_offset

    self.space_min = space_range[0]
    self.space_max = space_range[1]
    self.space_length = self.space_max - self.space_min
    self.space_offset_min = space_offset[0]
    self.space_offset_max = space_offset[1]

    self.dataset_size = dataset_size
    self.random_seed = random_seed
    self.rng = np.random.RandomState(random_seed)

    self.noise_density = noise_density
    self.pattern_density = pattern_density

    self.quantization_size = quantization_size
    self.dimension = dimension


class HyperLineDataset(HyperPlaneDataset):

  def __getitem__(self, i):
    # Regardless of the input index, return randomized data
    normal = self.rng.rand(self.dimension) - 0.5 + self.normal_offset
    normal /= np.sqrt((normal**2).sum())
    intercept = self.intercept_slope * self.rng.rand(
        self.dimension) + self.intercept_bias

    # Sample points on uniformly on space_range
    n_line = int(
        np.floor(self.pattern_density * np.sqrt(self.dimension) * self.space_length))
    X_line = self.space_length * (self.rng.rand(n_line, 1) -
                                  0.5) * normal[None, :] + intercept

    n_noise = int(np.floor(self.noise_density * (self.space_length**self.dimension)))
    X_noise = self.space_length * self.rng.rand(
        n_noise, self.dimension) + self.space_min + 0.1 * self.space_length * self.rng.randn(
            n_noise, self.dimension)

    space_offset = self.rng.rand(self.dimension) * (
        self.space_offset_max - self.space_offset_min) + self.space_offset_min

    # Concatenate data
    input = np.vstack((X_line, X_noise)) + space_offset[None, :]
    label = np.vstack([np.ones((len(X_line), 1)), np.zeros((len(X_noise), 1))])

    # Discretize
    discrete_coords = np.floor(input / self.quantization_size)
    keys = ME.utils.ravel_hash_vec(discrete_coords.copy())
    _, inds, counts = np.unique(keys, return_index=True, return_counts=True)
    return discrete_coords[inds], input[inds], counts[:, None], label[
        inds], normal, intercept